MaxPool op resolved as Aten OP

I created a small model training script, to test out maxpool gradient op for OneDNN EP, and the model definition below was used. But for some reason, the maxpool was resolved to Aten Op in the onnx graph. Is there a way to force torch-ort to use maxpool instead aten op? (maybe by disabling use of aten ops?)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.maxpool1 = nn.MaxPool2d(2)
        self.maxpool2 = nn.MaxPool2d(2)
        #self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(self.maxpool1(self.conv1(x)))
        x = F.relu(self.maxpool2(self.conv2(x)))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

[baijumeswani]

@chethanpk I just tried this on my end as well. If I comment out this code in my local pip directory, I can avoid generating the Aten nodes that you’re seeing in your graph. Here is the sample I am using (from you example):

import torch.nn as nn
import torch.nn.functional as F
import torch
from onnxruntime.training import ORTModule

from onnxruntime.training.ortmodule import DebugOptions

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.maxpool1 = nn.MaxPool2d(2)
        self.maxpool2 = nn.MaxPool2d(2)
        #self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(self.maxpool1(self.conv1(x)))
        x = F.relu(self.maxpool2(self.conv2(x)))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

x = torch.randn([1, 1, 30, 30], dtype=torch.float32)
net = ORTModule(Net(), DebugOptions(save_onnx=True, onnx_prefix='test'))
out = net(x)
print(out)

Here is the graph I get:

[chethanpk]

Rebuilding the wheel worked, thanks.