Remove .clone() calls in profiler

ghstack-source-id: fb9e0c0815328c2feacdbf4dec8164d52e0e951e
Pull Request resolved: https://github.com/fairinternal/xformers/pull/485

__original_commit__ = fairinternal/xformers@6f3f928d8c709d05d6dd353c4a42128b43744719

xformers/profiler/profiler.py

@@ -17,7 +17,6 @@
 import torch.nn as nn
 import torch.profiler
 import torch.utils.hooks
-from torch.utils._pytree import tree_map
 
 logger = logging.getLogger(__name__)
 
@@ -267,9 +266,6 @@ def _enter_module_hook(self, name):
         class PopState(torch.autograd.Function):
             @staticmethod
             def forward(ctx, *args):
-                args = tree_map(
-                    lambda x: x.clone() if isinstance(x, torch.Tensor) else x, args
-                )
                 if len(args) == 1:
                     return args[0]
                 return args
@@ -291,9 +287,6 @@ def _exit_module_hook(self, name):
         class PushState(torch.autograd.Function):
             @staticmethod
             def forward(ctx, *args):
-                args = tree_map(
-                    lambda x: x.clone() if isinstance(x, torch.Tensor) else x, args
-                )
                 if len(args) == 1:
                     return args[0]
                 return args