Hi I have a transformer encoder and I’m testing it with the following code:
import torch
import torch.autograd.profiler as profiler

encoder = torch.jit.load('eval/encoder.zip')


tmp = torch.ones([1, 7, 80])
len = torch.Tensor([7])

#Warmup
encoder.forward(tmp, len)
encoder.forward(tmp, len)

print("PROFILING ONE SHOT ENCODE")
with profiler.profile(with_stack=True, profile_memory=True) as prof:
    # Input is of the form [batch_size, len, feature_depth]
    out, _, _ = encoder.forward(torch.ones([1, 403, 80]), torch.Tensor([403]))
    print("DONE WITH FORWARD")
    print(prof.key_averages(group_by_stack_n=15).table(sort_by='self_cpu_time_total'))  


And I’m getting the following error :
RuntimeError: can't export a trace that didn't finish running

Even though I see that DONE message (which I presume means the encoder forward is done). Does anyone know why this might be happening?
The encoder code itself is pretty similar to fairseq/vggtransformer.py at master · pytorch/fairseq · GitHub but unfortunately I can’t give you the jit file or the exact code