PyTorch inferencer from scratch using trained weights from NVidia Modulus Wave_1D example

I am trying to create an inferencer that will initialize from a .pth trained model file. This is the wave_1d example from the NVidia Modulus examples. I trained it using the containerized modulus, but I want to use the trained model elsewhere. I am early on in my learning here but I have this python code using pytorch. It loads the weights and biases as I expect, and I can see the values match what was in the file, which seems to be an OrderedDict. You can see how I set up my neural net class, and it runs, but I do not get the answers I would expect. Any help would be appreciated to let me know what I am missing.
==============================================>8==============================================

import numpy as np
import torch as pt
import torch.nn as nn
import torch.nn.functional as f
import torch.optim as optim
import collections as coll
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt

pt.set_grad_enabled (False) 

class Wave1D(nn.Module):
    def __init__(self, numinputs, numoutputs):
        super(Wave1D, self).__init__()
        self.layer0 = nn.utils.weight_norm(nn.Linear(numinputs, 512),name='weight', dim=0).cuda()
        self.layer1 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer2 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer3 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer4 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer5 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer6 = nn.Linear(512, numoutputs).cuda()
        self.layer0.eval()
        self.layer1.eval()
        self.layer2.eval()
        self.layer3.eval()
        self.layer4.eval()
        self.layer5.eval()
        self.layer6.eval()
        
    def forward(self, x):
        x0 = f.silu(self.layer0(x))
        x1 = f.silu(self.layer1(x0))
        x2 = f.silu(self.layer2(x1))
        x3 = f.silu(self.layer3(x2))
        x4 = f.silu(self.layer4(x3))
        x5 = f.silu(self.layer5(x4))
        return f.silu(self.layer6(x5))
    
    def LoadFromOrderedDict(self, od):
        for k, v in od.items():
            if(k == '_impl.layers.0.linear.weight'): 
                print('Loading: ',k, np.shape(v), ' -> self.layer0.weight.data')
                self.layer0.weight_v.data = v
                self.layer0.weight_v.requires_grad = False
            if(k == '_impl.layers.0.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer0.weight_g.data')
                self.layer0.weight_g.data = v
                self.layer0.weight_g.requires_grad = False
            if(k == '_impl.layers.0.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer0.bias.data')
                self.layer0.bias.data = v
                self.layer0.bias.requires_grad = False
            if(k == '_impl.layers.1.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer1.weight.data')
                self.layer1.weight_v.data = v
                self.layer1.weight_v.requires_grad = False
            if(k == '_impl.layers.1.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer1.weight_g.data')
                self.layer1.weight_g.data = v
                self.layer1.weight_g.requires_grad = False
            if(k == '_impl.layers.1.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer1.bias.data')
                self.layer1.bias.data = v
                self.layer1.weight_g.requires_grad = False
            if(k == '_impl.layers.2.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer2.weight.data')
                self.layer2.weight_v.data = v
                self.layer2.weight_v.requires_grad = False
            if(k == '_impl.layers.2.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer2.weight_g.data')
                self.layer2.weight_g.data = v
                self.layer2.weight_g.requires_grad = False
            if(k == '_impl.layers.2.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer2.bias.data')
                self.layer2.bias.data = v
                self.layer2.bias.requires_grad = False
            if(k == '_impl.layers.3.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer3.weight.data')
                self.layer3.weight_v.data = v
                self.layer3.weight_v.requires_grad = False
            if(k == '_impl.layers.3.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer3.weight_g.data')
                self.layer3.weight_g.data = v
                self.layer3.weight_g.requires_grad = False
            if(k == '_impl.layers.3.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer3.bias.data')
                self.layer3.bias.data = v
                self.layer3.bias.requires_grad = False
            if(k == '_impl.layers.4.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer4.weight.data')
                self.layer4.weight_v.data = v
                self.layer4.weight_v.requires_grad = False
            if(k == '_impl.layers.4.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer4.weight_g.data')
                self.layer4.weight_g.data = v
                self.layer4.weight_g.requires_grad = False
            if(k == '_impl.layers.4.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer4.bias.data')
                self.layer4.bias.data = v
                self.layer4.bias.requires_grad = False
            if(k == '_impl.layers.5.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer5.weight.data')
                self.layer5.weight_v.data = v
                self.layer5.weight_v.requires_grad = False
            if(k == '_impl.layers.5.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer5.weight_g.data')
                self.layer5.weight_g.data = v
                self.layer5.weight_g.requires_grad = False
            if(k == '_impl.layers.5.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer5.bias.data')
                self.layer5.bias.data = v
                self.layer5.bias.requires_grad = False
            if(k == '_impl.final_layer.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer6.weight.data')
                self.layer6.weight.data = v
                self.layer6.weight.requires_grad = False
            if(k == '_impl.final_layer.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer6.bias.data')
                self.layer6.bias.data = v
                self.layer6.bias.requires_grad = False

    def dumpLayerWeightSizes(self):
        print(self.layer0.weight.size())
        print(self.layer1.weight.size())
        print(self.layer2.weight.size())
        print(self.layer3.weight.size())
        print(self.layer4.weight.size())
        print(self.layer5.weight.size())
        print(self.layer6.weight.size())

    def dumpLayerWeights(self):
        print(self.layer0.weight)
        print(self.layer1.weight)
        print(self.layer2.weight)
        print(self.layer3.weight)
        print(self.layer4.weight)
        print(self.layer5.weight)
        print(self.layer6.weight)

    def dumpLayerWeightGSizes(self):
        print(self.layer0.weight_g.size())
        print(self.layer1.weight_g.size())
        print(self.layer2.weight_g.size())
        print(self.layer3.weight_g.size())
        print(self.layer4.weight_g.size())
        print(self.layer5.weight_g.size())

    def dumpLayerWeightGs(self):
        print(self.layer0.weight_g)
        print(self.layer1.weight_g)
        print(self.layer2.weight_g)
        print(self.layer3.weight_g)
        print(self.layer4.weight_g)
        print(self.layer5.weight_g)

    def dumpLayerBiasSizes(self):
        print(self.layer0.bias.size())
        print(self.layer1.bias.size())
        print(self.layer2.bias.size())
        print(self.layer3.bias.size())
        print(self.layer4.bias.size())
        print(self.layer5.bias.size())
        print(self.layer6.bias.size())

    def dumpLayerBiases(self):
        print(self.layer0.bias)
        print(self.layer1.bias)
        print(self.layer2.bias)
        print(self.layer3.bias)
        print(self.layer4.bias)
        print(self.layer5.bias)
        print(self.layer6.bias)

model = Wave1D(2, 1)
model.eval()
originalOD = pt.load('C:\\\\Users\\IMSIR\\Documents\\ModulusTrainedModels\\wave_network.pth')
model.LoadFromOrderedDict(originalOD)

print('Original OrderedDict:')
for k, v in originalOD.items():
    print(k)
#print(originalOD)
print('=============================================================================================')
print('from model:')
print('SIZES:')
print('Weight Sizes:')
model.dumpLayerWeightSizes()
print('Weight_G Sizes:')
model.dumpLayerWeightGSizes()
print('Bias Sizes:')
model.dumpLayerBiasSizes()
print('WEIGHTS:')
model.dumpLayerWeights()
print('WEIGHT_GS:')
model.dumpLayerWeightGs()
print('BIASES:')
model.dumpLayerBiases()

t = 0.0
my2dspace = pt.tensor(np.linspace([0.0, t], [3.1415926535, t], num=101), requires_grad=False)
for i in range(1,101):
    t = t + 0.01*2*np.pi
    my2dspace = pt.cat((my2dspace, pt.tensor(np.linspace([0.0, t], [3.1415926535, t], num=101), requires_grad=False)), 0)

model.eval()

myOutput = model(my2dspace.float().cuda())
myCPUOutput = myOutput.cpu()

print('x:', my2dspace.select(1,0).float())
print('t:', my2dspace.select(1,1).float())
print('u:', myCPUOutput.squeeze().detach().numpy())
fig = plt.figure()
ax = plt.axes(projection ='3d')
ax.scatter(my2dspace.select(1,0).squeeze(), my2dspace.select(1,1).squeeze(), myCPUOutput.squeeze().detach().numpy())
plt.show()

del model
del my2dspace
del myOutput
del myCPUOutput
del originalOD

OK, I think I have cracked the code. Here is my PyTorch-based code to read the .pth file and it looks like it produces the correct result. I am a bit new to Python, so it is really about figuring out the shapes and converting between data types. Python has soooooo much stuff in it, no matter what you want to do there is a function in python to do it, the trick is finding it and then understanding what it does under the hood.

#read in the wave_network.pth file from the trained modulus wave_1d example, use it to inference a sampled space, and plot it.
import numpy as np
import torch as pt
import torch.nn as nn
import torch.nn.functional as f
import torch.optim as optim
import collections as coll
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
from matplotlib import cm

pt.set_grad_enabled (False) 

class Wave1D(nn.Module):
    def __init__(self, numinputs, numoutputs):
        super(Wave1D, self).__init__()
        self.layer0 = nn.utils.weight_norm(nn.Linear(numinputs, 512),name='weight', dim=0).cuda()
        self.layer1 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer2 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer3 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer4 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer5 = nn.utils.weight_norm(nn.Linear(512, 512),name='weight', dim=0).cuda()
        self.layer6 = nn.Linear(512, numoutputs).cuda()
        self.layer0.eval()
        self.layer1.eval()
        self.layer2.eval()
        self.layer3.eval()
        self.layer4.eval()
        self.layer5.eval()
        self.layer6.eval()
        
    def forward(self, x):
        x0 = f.silu(self.layer0(x))
        x1 = f.silu(self.layer1(x0))
        x2 = f.silu(self.layer2(x1))
        x3 = f.silu(self.layer3(x2))
        x4 = f.silu(self.layer4(x3))
        x5 = f.silu(self.layer5(x4))
        return self.layer6(x5)
    
    def LoadFromOrderedDict(self, od):
        for k, v in od.items():
            if(k == '_impl.layers.0.linear.weight'): 
                print('Loading: ',k, np.shape(v), ' -> self.layer0.weight.data')
                self.layer0.weight_v.data = v
                self.layer0.weight_v.requires_grad = False
            if(k == '_impl.layers.0.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer0.weight_g.data')
                self.layer0.weight_g.data = v
                self.layer0.weight_g.requires_grad = False
            if(k == '_impl.layers.0.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer0.bias.data')
                self.layer0.bias.data = v
                self.layer0.bias.requires_grad = False
            if(k == '_impl.layers.1.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer1.weight.data')
                self.layer1.weight_v.data = v
                self.layer1.weight_v.requires_grad = False
            if(k == '_impl.layers.1.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer1.weight_g.data')
                self.layer1.weight_g.data = v
                self.layer1.weight_g.requires_grad = False
            if(k == '_impl.layers.1.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer1.bias.data')
                self.layer1.bias.data = v
                self.layer1.weight_g.requires_grad = False
            if(k == '_impl.layers.2.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer2.weight.data')
                self.layer2.weight_v.data = v
                self.layer2.weight_v.requires_grad = False
            if(k == '_impl.layers.2.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer2.weight_g.data')
                self.layer2.weight_g.data = v
                self.layer2.weight_g.requires_grad = False
            if(k == '_impl.layers.2.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer2.bias.data')
                self.layer2.bias.data = v
                self.layer2.bias.requires_grad = False
            if(k == '_impl.layers.3.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer3.weight.data')
                self.layer3.weight_v.data = v
                self.layer3.weight_v.requires_grad = False
            if(k == '_impl.layers.3.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer3.weight_g.data')
                self.layer3.weight_g.data = v
                self.layer3.weight_g.requires_grad = False
            if(k == '_impl.layers.3.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer3.bias.data')
                self.layer3.bias.data = v
                self.layer3.bias.requires_grad = False
            if(k == '_impl.layers.4.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer4.weight.data')
                self.layer4.weight_v.data = v
                self.layer4.weight_v.requires_grad = False
            if(k == '_impl.layers.4.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer4.weight_g.data')
                self.layer4.weight_g.data = v
                self.layer4.weight_g.requires_grad = False
            if(k == '_impl.layers.4.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer4.bias.data')
                self.layer4.bias.data = v
                self.layer4.bias.requires_grad = False
            if(k == '_impl.layers.5.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer5.weight.data')
                self.layer5.weight_v.data = v
                self.layer5.weight_v.requires_grad = False
            if(k == '_impl.layers.5.linear.weight_g'):
                print('Loading: ',k, np.shape(v), ' -> self.layer5.weight_g.data')
                self.layer5.weight_g.data = v
                self.layer5.weight_g.requires_grad = False
            if(k == '_impl.layers.5.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer5.bias.data')
                self.layer5.bias.data = v
                self.layer5.bias.requires_grad = False
            if(k == '_impl.final_layer.linear.weight'):
                print('Loading: ',k, np.shape(v), ' -> self.layer6.weight.data')
                self.layer6.weight.data = v
                self.layer6.weight.requires_grad = False
            if(k == '_impl.final_layer.linear.bias'):
                print('Loading: ',k, np.shape(v), ' -> self.layer6.bias.data')
                self.layer6.bias.data = v
                self.layer6.bias.requires_grad = False

    def dumpLayerWeightSizes(self):
        print(self.layer0.weight.size())
        print(self.layer1.weight.size())
        print(self.layer2.weight.size())
        print(self.layer3.weight.size())
        print(self.layer4.weight.size())
        print(self.layer5.weight.size())
        print(self.layer6.weight.size())

    def dumpLayerWeights(self):
        print(self.layer0.weight)
        print(self.layer1.weight)
        print(self.layer2.weight)
        print(self.layer3.weight)
        print(self.layer4.weight)
        print(self.layer5.weight)
        print(self.layer6.weight)

    def dumpLayerWeightGSizes(self):
        print(self.layer0.weight_g.size())
        print(self.layer1.weight_g.size())
        print(self.layer2.weight_g.size())
        print(self.layer3.weight_g.size())
        print(self.layer4.weight_g.size())
        print(self.layer5.weight_g.size())

    def dumpLayerWeightGs(self):
        print(self.layer0.weight_g)
        print(self.layer1.weight_g)
        print(self.layer2.weight_g)
        print(self.layer3.weight_g)
        print(self.layer4.weight_g)
        print(self.layer5.weight_g)

    def dumpLayerBiasSizes(self):
        print(self.layer0.bias.size())
        print(self.layer1.bias.size())
        print(self.layer2.bias.size())
        print(self.layer3.bias.size())
        print(self.layer4.bias.size())
        print(self.layer5.bias.size())
        print(self.layer6.bias.size())

    def dumpLayerBiases(self):
        print(self.layer0.bias)
        print(self.layer1.bias)
        print(self.layer2.bias)
        print(self.layer3.bias)
        print(self.layer4.bias)
        print(self.layer5.bias)
        print(self.layer6.bias)

model = Wave1D(2, 1)
model.eval()
originalOD = pt.load('C:\\\\Users\\IMSIR\\Documents\\ModulusTrainedModels\\wave_network.pth')
model.LoadFromOrderedDict(originalOD)

print('Original OrderedDict:')
for k, v in originalOD.items():
    print(k)
#print(originalOD)
print('=============================================================================================')
print('from model:')
print('SIZES:')
print('Weight Sizes:')
model.dumpLayerWeightSizes()
print('Weight_G Sizes:')
model.dumpLayerWeightGSizes()
print('Bias Sizes:')
model.dumpLayerBiasSizes()
print('WEIGHTS:')
model.dumpLayerWeights()
print('WEIGHT_GS:')
model.dumpLayerWeightGs()
print('BIASES:')
model.dumpLayerBiases()

t = 0.0
my2dspace = pt.tensor(np.linspace([0.0, t], [3.141592653589793, t], num=313), requires_grad=False)
for i in range(1,628):
    t = t + 0.01
    my2dspace = pt.cat((my2dspace, pt.tensor(np.linspace([0.0, t], [3.141592653589793, t], num=313), requires_grad=False)), 0)

model.eval()

myOutput = model(my2dspace.float().cuda())
myCPUOutput = myOutput.cpu()

xx = np.linspace(0.0, 3.141592653589793, num=313)
tt = np.linspace(0.0, 2.0*3.141592653589793, num=628)
xx, tt = np.meshgrid( xx, tt )
uu = myCPUOutput.numpy().reshape(628, 313)
print('xx: ', xx)
print('tt: ', tt)
print('uu: ', uu)

myCPUOutput.squeeze().detach().numpy())
#surf = ax.plot_surface(xx, tt, uu, cmap=cm.coolwarm, linewidth=5, antialiased=False)
fig, ax = plt.subplots()
surf = ax.pcolormesh(xx, tt, uu, vmin=-1.0, vmax=1.0)
plt.show()

del model
del my2dspace
del myOutput
del myCPUOutput
del originalOD