Not sure if I’m right but I think you can specify the device to load the traced/saved model to in the load function with the  map_location parameter
https://pytorch.org/docs/master/generated/torch.jit.load.html#torch.jit.load 94