[hotfix]fixed bugs of assigning grad states to non leaf nodes (#711)

* fixed bugs of assigning grad states to non leaf nodes * use detach()
3 years ago · 140263a394
1 changed files with 1 additions and 1 deletions
--- a/colossalai/utils/activation_checkpoint.py
+++ b/colossalai/utils/activation_checkpoint.py
@ -12,7 +12,7 @@ def copy_to_device(obj, device):
    if torch.is_tensor(obj):
        # Notice:
        # When in no_grad context, requires_gard is False after movement
-        ret = obj.to(device)
+        ret = obj.to(device).detach()
        ret.requires_grad = obj.requires_grad
        return ret
    elif isinstance(obj, list):