return mean.sub(std2), mean.add(std2)
def cpu(self):
representation = self.representation()
if not isinstance(representation, tuple) or isinstance(representation, list):
representation = representation,
return self.__class__(*(var.cpu() for var in representation))
After Change
return mean.sub(std2), mean.add(std2)
def cpu(self):
new_args = []
new_kwargs = {}
for arg in self._args:
if hasattr(arg, "cpu"):
new_args.append(arg.cpu())
else:
new_args.append(arg)
for name, val in self._kwargs.items():
if hasattr(val, "cpu"):