TensorRT 7 conv3d is not running on Tensor Cores

AI & Data Science Deep Learning (Training & Inference) TensorRT
8

Hi,

I’ve double checked and replaced in my model creation code Conv3d/MaxPool3d with Conv2d/MaxPool2d and i see Conv2d executing on Tensor Cores.

import torch
from torch import nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(64, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.conv(x)
        x = F.relu(x)
        return x

#image_dims = (1, 64, 64, 64, 64)
image_dims = (1, 64, 64, 64)

dummy_input = torch.rand(image_dims, device="cuda")

model = Net().to("cuda")

torch.onnx.export(model, dummy_input, "test.onnx", verbose=True, opset_version=11)

Can you please share a link or some documentation on cudNN 8.0 for jetson platforms, because as I got it, it differs from desktop release and it’s not clear, what is supported and what is not?