tensor
Utility functions for PyTorch tensors.
Functions
Convert numpy arrays to torch tensors. |
|
Return a context manager that sets the CUDA device to be the same as the input tensor. |
|
Move tensors to device if not meta device; otherwise materialize with empty_like(). |
|
Try to recursively detach the data from the computation graph. |
|
Try to recursively move the data to the specified args/kwargs. |
|
Convert torch tensors to numpy arrays. |
- numpy_to_torch(np_outputs)
Convert numpy arrays to torch tensors.
- Parameters:
np_outputs (list[ndarray])
- Return type:
list[Tensor]
- same_device_as(inputs)
Return a context manager that sets the CUDA device to be the same as the input tensor.
Returns a null context if the tensor is on CPU or on the same device as the current CUDA device.
- Parameters:
inputs (Tensor)
- to_empty_if_meta_device(module, *, device, recurse=True)
Move tensors to device if not meta device; otherwise materialize with empty_like().
Officially, torch suggests to_empty() for meta device materialization. Under the hood, torch.empty_like() is applied to all parameters or buffers (see _apply). This may accidentally overwrite buffers with precomputed values during construction. Given the goal is to only materialize those tensors on meta device, this function checks the device first and only move the tensor to the destination if it is not on meta device.
- Parameters:
module (Module) – The target module to apply this transformation.
device (device) – The desired device of the parameters and buffers in this module.
recurse – Whether parameters and buffers of submodules should be recursively moved to the specified device.
- torch_detach(data)
Try to recursively detach the data from the computation graph.
- torch_to(data, *args, **kwargs)
Try to recursively move the data to the specified args/kwargs.
- torch_to_numpy(inputs)
Convert torch tensors to numpy arrays.
- Parameters:
inputs (list[Tensor])
- Return type:
list[ndarray]