def METHOD_NAME(self, *args, to_cpu=False, **kw_args):
    with torch.no_grad():
        with self.amp_manager.context():
            D = super().METHOD_NAME(*args, **kw_args)
            if to_cpu:
                return (D[0].cpu(), *D[1:]) if isinstance(D, tuple) else D.cpu()
            return D