Cannot move ScriptModule to GPU with to()

(关鸿) #1

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.

(Michael Suo) #2

What version of PyTorch are you using? This should be available in 1.0 or the nightlies

(Justus Schock) #3

You can’t do this with .to but net.cuda() should work.

(关鸿) #4

net.cuda() works. Thank you. Do you think Pytorch should add .to() to ScriptModules?

(Justus Schock) #5

To answer your last question: @tom told me, that this is enabled on github master and nightlies. If you want that feature, you should probably switch to one of these.