Hi, I’ve been trying to migrate some code from torch 1.0.1 to torch 1.3.1 but I’m struggling with an error that appears only with the latest version.
Here is a minimal reproducible code to illustrate my problem (just a random MLP with embedding layers for first 3 categorical columns) :
import torch
class NN_emb_long(torch.nn.Module):
    def __init__(self,):
        super(NN_emb_long, self).__init__()
        emb_layers = []
        for i in range(3):
            emb_layers.append(torch.nn.Embedding(5, 2))
        self.emb_layers = torch.nn.ModuleList(emb_layers)        
        self.lin1 = torch.nn.Linear(8, 16)
        self.lin_out = torch.nn.Linear(16, 2)
    
    def forward(self, x):
        embs_list = []
        for i in range(3):
            embs = self.emb_layers[i](x[:,i].long())
            embs_list.append(embs)
        post_embed = torch.cat([x[:, torch.Tensor([3,4]).long()]]+embs_list, dim=1)
        res = self.lin1(post_embed)
        res = torch.nn.ReLU()(res)
        res = self.lin_out(res)
        return res, post_embed

NN = NN_emb_long()
input_example = torch.ones((10, 5)).requires_grad_(True)

probas, post_embeddings = NN(input_example)
grad_outputs = torch.ones(input_example.shape[0],2)
G = torch.autograd.grad(outputs=probas,
                        inputs=post_embeddings,
                        grad_outputs=grad_outputs,
                        only_inputs=True,
                        retain_graph=True
                       )[0]
print(G.shape)
# until this everything should work on both version
# taking the trace shows an error only with torch 1.3.1

basic_trace = torch.jit.trace(NN, input_example, check_trace=True)
probas,  post_embeddings = basic_trace(input_example)
grad_outputs = torch.ones(10,2)
G = torch.autograd.grad(outputs=probas,
                        inputs=post_embeddings,
                        grad_outputs=grad_outputs,
                        only_inputs=True,
                        retain_graph=True,
                       )[0]
print(G.shape)


This code should run fine with torch 1.0.1 but fail with the following error with torch 1.3.1:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-11-d55fb440d76f> in <module>
      6                         grad_outputs=grad_outputs,
      7                         only_inputs=True,
----> 8                         retain_graph=True,
      9                        )[0]
     10 print(G.shape)

.cache/poetry/engine-py3.6/lib/python3.6/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    155     return Variable._execution_engine.run_backward(
    156         outputs, grad_outputs, retain_graph, create_graph,
--> 157         inputs, allow_unused)
    158 
    159 

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

Could you please help me understand what happens here?
Thank you!