pytorch
da322ea8 - Enable torch.jit.load for custom device (#99535)

Commit
1 year ago
Enable torch.jit.load for custom device (#99535) Fixes #ISSUE_NUMBER 1、torch.jit.load for custom device ``` # custom device named `foo` ts_model = torch.jit.script(mode.to(device="foo")) ts_model.save("./ts.pt") # it is a script model on device `foo` # and then we want to load it and run it torch.jit.load("./ts.pt") ``` 2、 add some extra key for custom device with `privateuse1` Pull Request resolved: https://github.com/pytorch/pytorch/pull/99535 Approved by: https://github.com/albanD
Committer
Parents
Loading