pytorch_pfn_extras.get_xp

pytorch_pfn_extras.get_xp(obj)

Returns a module of ndarray implementation (numpy or cupy) for the given obj.

The obj can be torch.Tensor, torch.device or NumPy/CuPy ndarray.

Parameters

obj (Union[Any, torch.Tensor]) –

Return type

Any

pytorch_pfn_extras.as_ndarray pytorch_pfn_extras.as_numpy_dtype