Metadata-Version: 2.4
Name: unify-tensors
Version: 0.1.0
Summary: A utility to unify the device of multiple PyTorch tensors.
Author-email: Giacomo Guiduzzi <guiduzzi.giacomo@gmail.com>
Project-URL: Homepage, https://github.com/giacomoguiduzzi/unify-tensors
Project-URL: Repository, https://github.com/giacomoguiduzzi/unify-tensors
Project-URL: Issues, https://github.com/giacomoguiduzzi/unify-tensors/issues
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Dynamic: license-file

# Unify Tensors
Tiny package to implement a safety function that moves all the tensors passed to the same device. 
Useful when a part of a big model has been offloaded and there's a risk of a runtime error because of tensors on different devices.
## Installation

```bash
pip install unify-tensors
```

or, if not available on PyPI, clone the repository and install it manually:

```bash
git clone https://github.com/giacomoguiduzzi/unify-tensors.git
cd unify-tensors
pip install .
```

## Usage
The `unify_tensors` function takes any number of tensor arguments and moves them all to the device of the first tensor provided.
If you want to move them to a specific device, you can pass the `device` keyword argument.

```python
from unify_tensors import unify_tensors
import torch

a = torch.randn(2, 2)
b = torch.randn(2, 2).to('cuda')
c = torch.randn(2, 2).to('cpu')

a, b, c = unify_tensors(a, b, c)

# Should print 'cpu' because the first tensor is on CPU
print(a.device, b.device, c.device)

a = a.to('cuda')
a, b, c = unify_tensors(a, b, c)

# Should print 'cuda' because the first tensor is on CUDA
print(a.device, b.device, c.device)

# You can also specify a device explicitly
a, b, c = unify_tensors(a, b, c, device='cpu')

# Should print 'cpu' because we specified the device
print(a.device, b.device, c.device)
```

## License
This project is licensed under the MIT License.
See the [LICENSE](LICENSE) file for details.
