Fix type hints (#45)

Summary: Pull Request resolved: #45 Reviewed By: anana10c Differential Revision: D65882392 Pulled By: tsunghsienlee fbshipit-source-id: f57a74be3c4bdaf1fdda869f30606b1f80d2a76d
facebookresearch · Nov 13, 2024 · 7ffbf1b · 7ffbf1b
1 parent 13f68b7
commit 7ffbf1b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/distributed_shampoo/utils/shampoo_ddp_distributor.py b/distributed_shampoo/utils/shampoo_ddp_distributor.py
@@ -78,9 +78,9 @@ def __init__(
 
         # Determine communication type.
         if distributed_config.communication_dtype == CommunicationDType.BF16:
-            communication_dtype: torch.dtype = torch.bfloat16
+            communication_dtype = torch.bfloat16
         elif distributed_config.communication_dtype == CommunicationDType.FP16:
-            communication_dtype: torch.dtype = torch.float16
+            communication_dtype = torch.float16
         else:
             assert distributed_config.communication_dtype in [
                 CommunicationDType.FP32,

diff --git a/distributed_shampoo/utils/shampoo_hsdp_distributor.py b/distributed_shampoo/utils/shampoo_hsdp_distributor.py
@@ -152,9 +152,9 @@ def __init__(
 
         # Determine communication type.
         if distributed_config.communication_dtype == CommunicationDType.BF16:
-            communication_dtype: torch.dtype = torch.bfloat16
+            communication_dtype = torch.bfloat16
         elif distributed_config.communication_dtype == CommunicationDType.FP16:
-            communication_dtype: torch.dtype = torch.float16
+            communication_dtype = torch.float16
         else:
             assert distributed_config.communication_dtype in [
                 CommunicationDType.FP32,