I use a simple network.
class Net(torch.nn.Module):
    def __init__(self, D_in=24, H=[128,128,128,128], D_out=1, drate = 0.0):
        
        super(Net, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H[0], bias= True)
        self.linear2 = torch.nn.Linear(H[0], H[1], bias= True)
        self.linear3 = torch.nn.Linear(H[1], H[2], bias= True)
        self.linear4 = torch.nn.Linear(H[2], H[3], bias= True)
        self.linear5 = torch.nn.Linear(H[3], D_out, bias = True)
        self.dlayer = torch.nn.Dropout(p=drate)
    
    def forward(self, x):
        x = self.dlayer(x)
        x = F.relu(self.linear1(x))
        x = self.dlayer(x)
        x = F.relu(self.linear2(x))
        x = self.dlayer(x)
        x = F.relu(self.linear3(x))
        x = self.dlayer(x)
        x = F.relu(self.linear4(x))
        x = self.dlayer(x)
        x = self.linear5(x)

        return x

The cudnn verison is 8.04 (torch.backends.cudnn.version() gives 8004).
I used Adam to train the network.