Thank you, I managed to get it working in a simple case where the network takes one variable as input, but when I pass a list of variables, I get an error:
TypeError: _create_method_from_trace(): incompatible function arguments. The following argument types are supported:
1. (self: torch._C.ScriptModule, arg0: unicode, arg1: function, arg2: List[torch::autograd::Variable]) -> None
Here is a simple example that does produces this error:
import torch
import torch.jit
import torch.nn as nn
from torch.autograd import Variable
class TestNet(nn.Module):
def init(self):
super(TestNet, self).init()
self.net0 = nn.Linear(100, 200)
self.net1 = nn.Linear(100, 200)
self.net2 = nn.Linear(200, 1)
self.bn = nn.BatchNorm1d(200)
self.sigmoid = nn.Sigmoid()
self.ReLU = nn.ReLU(inplace=False)
self.drop = nn.Dropout(0.5)
…
def forward(self, V):
return self.sigmoid(self.net2(self.drop(self.ReLU(self.bn(self.net0(V[0])+self.net1(V[1])))))).squeez
use_cuda = True
net = TestNet()
if use_cuda:
net.cuda()
criterion.cuda()
V = Variable(torch.randn(100, 100)).cuda()
label = Variable(torch.randn(100)).cuda()
else:
V = Variable(torch.randn(100, 100))
label = Variable(torch.randn(100))
V = [V,V*2]
net.train()
net = torch.jit.trace(V)(net)
optim_betas = (0.9, 0.999)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.01, betas=optim_betas)
criterion = nn.BCELoss()
for i in range(0,1000000):
r = net(V)
Is this a bug or am I doing something wrong here?
Thank you