Using optimize_for_inference on Torchscript model causes CUDNN_BACKEND_OPERATIONGRAPH_DESCRIPTOR error

Question

I fine-tuned the PyTorch I3D action detection model on a custom dataset, saved it to Torch script and I’m loading it for inference by:

# the model is
model = torch.hub.load("facebookresearch/pytorchvideo", "i3d_r50", pretrained=True)

# training on custom dataset
[...]

# save the trained model
torch.jit.save(torch.jit.script(model), some_local_path)

# load it back
model = torch.jit.load(some_local_path, map_location=device).to(device)

# run inference
predictions = model(some_input)

If I just use the model loaded this way everything is ok...however, if I try adding:

model = torch.jit.load(some_local_path, map_location=device).to(device)
model = torch.jit.optimize_for_inference(model)
predictions = model(some_input)

I'm getting the error:

Traceback of TorchScript (most recent call last):

    graph(%input, %weight, %bias, %stride:int[], %padding:int[], %dilation:int[], %groups:int):
        %res = aten::cudnn_convolution_relu(%input, %weight, %bias, %stride, %padding, %dilation, %groups)
               ~~~~ <--- HERE
        return (%res)
RuntimeError: CUDNN_BACKEND_OPERATIONGRAPH_DESCRIPTOR: cudnnFinalize Failed cudnn_status: CUDNN_STATUS_NOT_SUPPORTED

Is this a problem in the optimization of the specific model, something due to the CUDA version?