Fix uninitialized mixed_D in dense class

pytorch · Nov 25, 2024 · 8cbcdff · 8cbcdff
1 parent 08a71b6
commit 8cbcdff
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py
@@ -3436,6 +3436,14 @@ def __init__(
             torch.tensor(D_offsets, device=self.current_device, dtype=torch.int32),
         )
         assert self.D_offsets.numel() == T + 1
+
+        mixed_D = False
+        D = dims[0]
+        for d in dims:
+            if d != D:
+                mixed_D = True
+                break
+        self.mixed_D = mixed_D
         # Required for VBE
         self.register_buffer(
             "feature_dims",