update test case to take into account param requirements and blocksize

nighting0le01 · Nov 7, 2024 · 668a4ba · 668a4ba
1 parent a4e320c
commit 668a4ba
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
@@ -390,10 +390,11 @@ def test_optim_exclude_low_bit_params(self, optim_name, dtype, device):
         self.assertTrue(exp_avg.__class__ == torch.Tensor)
         self.assertTrue(exp_avg_sq.__class__ == torch.Tensor)
         for param in model.parameters():
-            if id(param) not in excluded_params_ids :
+            if id(param) not in excluded_params_ids and param.numel() >= 4096 and param.numel() % optim.block_size == 0:
                 param_state = state[param]
                 exp_avg = param_state['exp_avg']
                 exp_avg_sq = param_state['exp_avg_sq']
+
                 self.assertTrue(exp_avg.__class__ != torch.Tensor)
                 self.assertTrue(exp_avg_sq.__class__ != torch.Tensor)