def METHOD_NAME(self, obj, device):
    if isinstance(obj, nn.Module):
        return obj.to(device)
    if torch.is_tensor(obj):
        return obj.to(device)
    if isinstance(obj, (tuple, list)):
        return [self.METHOD_NAME(el, device) for el in obj]
    if isinstance(obj, dict):
        return {
            name: self.METHOD_NAME(val, device)
            for name, val in obj.items()
        }
    raise ValueError(f'Unexpected type {type(obj)}')