def METHOD_NAME(self, device=None, non_blocking=False):
    if device is not None:
        return self.to(device=device, non_blocking=non_blocking)
    return self.to(device='cuda:0', non_blocking=non_blocking)