def __torch_function__(self, func, types, args=(), kwargs=None):
    if (func in _SHARDED_OPS):
        return _SHARDED_OPS[func](types, args, kwargs, self._process_group)
    raise RuntimeError(f"torch function '{func.__name__}', with args: {args} and kwargs: {kwargs} not supported for ShardedTensor!")
def __torch_function__(self, func, types, args=(), kwargs=None):
    if (func in _SHARDED_OPS):
        return _SHARDED_OPS[func](types, args, kwargs, self._process_group)
    raise RuntimeError(f"torch function '{func.__name__}', with args: {args} and kwargs: {kwargs} not supported for ShardedTensor!")