Add `map_location` when importing to Torch

neuromorphs / NIRTorch

PyTorch helper module to translate to and from NIR

https://neuroir.org/docs

BSD 3-Clause "New" or "Revised" License

7 stars 3 forks source link

Add `map_location` when importing to Torch #10

Open Jegp opened 1 year ago

Jegp commented 1 year ago

It's currently ambiguous which device imported modules and tensors are sent to. We should allow the user to specify something like the map_location parameter to disambiguate and simplify device acceleration.

https://pytorch.org/docs/stable/generated/torch.load.html