My code is
import torch
import torch.nn.functional as F
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class ScriptNet(torch.jit.ScriptModule):
def __init__(self, n_features, out_size):
super().__init__()
self.fc1 = torch.jit.trace(torch.nn.Linear(n_features, 50), torch.rand(1, n_features))
self.fc2 = torch.jit.trace(torch.nn.Linear(50, out_size), torch.rand(1, 50))
@torch.jit.script_method
def forward(self, x):
x = F.dropout(F.relu(self.fc1(x)), 0.5)
return F.softmax(self.fc2(x))
net = ScriptNet(10,2).to(device)
I got this error:
RuntimeError: to is not supported on TracedModules
I know I can save the ScriptModule and use torch.jit.load(modulefile, map_location=‘cuda’) to load the module into GPU, but I wonder if I can directly move the module to GPU without saving and loading.