Merge pull request #53 from Achazwl/fix-adam-torch12

fix adam API changed in torch>=1.12.0

Author: a710128

bmtrain/optim/adam.py

@@ -110,7 +110,8 @@ def step(self, closure=None, scale=1):
                             [state['exp_avg']],
                             [state["exp_avg_sq"]],
                             [],
-                            [state["step"]],
+                            [state["step"]] if int(torch.__version__.split('.')[1]) < 12
+                                else [torch.tensor(state["step"])],
                             amsgrad=False,
                             beta1=group['betas'][0],
                             beta2=group['betas'][1],

bmtrain/optim/adam_offload.py

@@ -136,7 +136,8 @@ def step(self, closure=None, scale=1):
                     [state["exp_avg"]],
                     [state["exp_avg_sq"]],
                     [],
-                    [state["step"]],
+                    [state["step"]] if int(torch.__version__.split('.')[1]) < 12
+                        else [torch.tensor(state["step"])],
                     amsgrad=False,
                     beta1=beta1,
                     beta2=beta2,

tests/test_optim.py

@@ -25,8 +25,8 @@ def main():
     model2.load_state_dict(state_dict)
     model3.load_state_dict(state_dict)
 
-    model1 = model1.cuda().half()
-    model2 = model2.cuda().half()
+    model1 = model1.cuda()
+    model2 = model2.cuda()
     model3 = model3.cuda()
 
     opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)