tiny style change

mosaicml · Aug 15, 2023 · b87ca31 · b87ca31
1 parent 476a9ec
commit b87ca31
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/optim/lion8b.py b/llmfoundry/optim/lion8b.py
@@ -306,7 +306,7 @@ def lion_step_unfused(grads: torch.Tensor,
                       beta2: float,
                       weight_decay: float = 0) -> torch.Tensor:
     # f32 cast to match fused impl + for compatibility with f32 grads or weights
-    momentums = momentums.to(torch.float32)
+    momentums = momentums.to(dtype=torch.float32)
     grads = grads.to(dtype=torch.float32)
 
     update = momentums.lerp(grads, 1 - beta1).sign_()