Merge pull request #8 from OpenBMB/FX_checkpointing_grad

FX: always requires_grad in checkpointing block
OpenBMB · Mar 16, 2022 · 4bde89a · 4bde89a
2 parents b5bc1ea + 32e49fb
commit 4bde89a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/bmtrain/block_layer.py b/bmtrain/block_layer.py
@@ -166,7 +166,10 @@ def enter(self):
             shape = param["shape"]
             param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape)
             if requires_grad and kw_name in self._grad_buffer:
+                param["parameter"].requires_grad_(True)
                 param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape)
+            else:
+                param["parameter"].requires_grad_(False)
 
 
     def __enter__(self):