JIT Script on Lazy param

[mprithiv]

I am trying to convert my GraphSAGE model to JIT script version based this example. I always keep gettingUninitialized PArameters Weight error for
LazyOAram

HeteroData( user={ num_nodes=305171, x=[305171, 1153], y=[305171], train_mask=[305171], val_mask=[305171], test_mask=[305171] }, seller={ num_nodes=31375, x=[31375, 771], train_mask=[31375], val_mask=[31375], test_mask=[31375] }, (user, amount, seller)={ edge_index=[2, 349379] }, (user, date_of_txn, seller)={ edge_index=[2, 349379] }, (seller, rev_amount, user)={ edge_index=[2, 349379] }, (seller, rev_date_of_txn, user)={ edge_index=[2, 349379] } )

My module looks like this

`class GNN(torch.nn.Module):
def init(self, hidden_channels, out_channels):
super().init()
torch.manual_seed(1234567)

    self.conv1 = SAGEConv((-1, -1), hidden_channels).jittable()
    self.conv2 = SAGEConv((-1, -1), out_channels).jittable()
    #self.classifier = Linear(out_channels, 2)


def forward(self, x, edge_index):
    x = self.conv1(x, edge_index)
    x = x.relu()
    x = F.dropout(x, p=0.5, training=self.training)
    x = self.conv2(x, edge_index)
    x = x.relu()
    #out = self.classifier(x)
    return x#out, x

model = GNN(hidden_channels=512, out_channels=32)
model = to_hetero(model, data.metadata(), aggr='sum')
model.to(device)
model = torch.jit.script(model)
`

The error was raise RuntimeError("'{}' has uninitialized parameters {}. Did you forget to run a forward pass?" .format(torch.typename(type(mod)), name))

[rusty1s]

I don't think TorchScript supports lazy parameters via -1. As such, you would first need to run a single forward pass of your model in order to initialize parameters, and then convert your model to TorchScript via:

model(x, edge_index)  # Run a forward pass
model.conv1 = model.conv1.jittable() 
model.conv2 = model.conv2.jittable()
model = torch.jit.script(model)