skip testing errors when not present

mosaicml · Oct 6, 2023 · 10aba20 · 10aba20
1 parent b24f095
commit 10aba20
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_lion8b.py b/tests/test_lion8b.py
@@ -481,7 +481,7 @@ def _set_state_dict_type(model: nn.Module):
 
         assert mom_orig.shape == mom_new.shape
         assert mom_orig.dtype == mom_new.dtype
-        if use_errors:
+        if use_errors and (dtype != torch.float32):
             errs_orig = d_orig['errors']
             errs_new = d_new['errors']
             assert errs_orig.shape == errs_new.shape