Nan error in VAE

PyTorch2603: <p>Hello. I have created the following VAE, but after the first few thousand epochs I get nan as the output when I print out the loss:<br>
import os<br>
import numpy as np<br>
import matplotlib.pyplot as plt<br>
import numpy as np<br>
import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F<br>
import torch.optim as optim<br>
import torchvision<br>
from torchvision import datasets, transforms<br>
from torch.autograd import Variable<br>
from torchvision.utils import save_image<br>
import os<br>
from PIL import Image as PImage<br>
device=“cpu”<br>
trainloss=0<br>
testloss=0<br>
def weights_init(m):<br>
if isinstance(m, nn.Linear):<br>
torch.nn.init.kaiming_uniform_(m.weight)<br>
class VAE(nn.Module):<br>
def <strong>init</strong>(self, z_dim):<br>
super(VAE, self).<strong>init</strong>()</p>
<pre><code>    # Encoder
    self.encode=nn.Sequential(
    nn.Linear(22,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720)
    )
    self.fc_mu = nn.Linear(720, z_dim)
    self.fc_logvar = nn.Linear(720, z_dim)
    # Decoder
    self.decode=nn.Sequential(
    nn.Linear(z_dim, 720),
    nn.ReLU(),
    nn.Linear(720, 720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.ReLU(),
    nn.Linear(720,720),
    nn.Tanh(),
    nn.Linear(720,22)
    )
    nn.init.kaiming_uniform_(self.fc_mu.weight)
    nn.init.kaiming_uniform_(self.fc_logvar.weight)
    self.encode.apply(weights_init)
    self.decode.apply(weights_init)

    
def encoder(self, x):
    a=self.encode(x)
    return self.fc_mu(a),self.fc_logvar(a)     

def sampling(self, mu, logvar):
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return eps.mul(std).add_(mu) # return z sample

def decoder(self, z):
    return self.decode(z)

def forward(self, x):
    mu, logvar = self.encoder(x)
    z = self.sampling(mu, logvar)
    return self.decoder(z), mu, logvar
</code></pre>
<p>vae = VAE(z_dim=20).to(device)<br>
optimizer = optim.Adam(vae.parameters(), lr=0.001)<br>
def loss_function(recon_x, x, mu, logvar):<br>
myloss=nn.MSELoss()<br>
MSE=myloss(recon_x,x)<br>
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())<br>
return MSE + KLD</p>
<p>def trainortest(myvector):<br>
optimizer.zero_grad()<br>
recon_batch, mu, logvar = vae(myvector)<br>
loss = loss_function(recon_batch, myvector, mu, logvar)<br>
if(epoch&lt;34000):<br>
loss.backward()<br>
myloss=loss.item()<br>
optimizer.step()<br>
else:<br>
myloss=loss.item()<br>
print(myloss)<br>
return myloss<br>
with open(‘x_y (1).txt’, ‘r’) as f:<br>
lines = f.readlines()<br>
epoch=0</p>
<p>trainloss=0<br>
testloss=0<br>
<span class="hashtag">#batch_size</span> will be 60<br>
for line in lines:<br>
epoch+=1<br>
print(epoch)<br>
x, y = line.split(‘=’)[0], line.split(‘=’)[1]<br>
x, y = x.split(’ ‘), y.split(’ ')<br>
x = [i for i in x if i]<br>
y = [i for i in y if i]<br>
x[0] = x[0][1:]<br>
y[0] = y[0][1:]<br>
x[-1] = x[-1][:-1]<br>
y[-1] = y[-1][:-1]</p>
<pre><code>    x = [float(i) for i in x if i]
    y = [float(i) for i in y if i]
    S=np.zeros(360, dtype='complex_')
    i=0
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=359
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=1
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=358
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=2
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=357
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=3
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=356
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=4
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=355
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    i=5
    for k in range(360):
        a=x[k]
        b=y[k]
        tmp = ((-2j*np.pi*i*k)) /360
        S[i] += (complex(a,b)) * np.exp(tmp)
    S[i]=S[i]/360
    input_list=[float(np.real(S[355])),float(np.real(S[356])),float(np.real(S[357])),float(np.real(S[358])), float(np.real(S[359])), float(np.real(S[0])), float(np.real(S[1])),float(np.real(S[2])),float(np.real(S[3])),float(np.real(S[4])),float(np.real(S[5])), float(np.imag(S[355])),float(np.imag(S[356])),float(np.imag(S[357])), float(np.imag(S[358])), float(np.imag(S[359])), float(np.imag(S[0])), float(np.imag(S[1])), float(np.imag(S[2])), float(np.imag(S[3])),float(np.imag(S[4])),float(np.imag(S[5]))]
    input_tensor=torch.tensor(input_list)
    answer=trainortest(input_tensor)

    if(epoch&lt;34000):
        trainloss+=answer
    else:
        testloss+=answer
</code></pre>
<p>print(trainloss/34000)<br>
print(testloss/(epoch-34000))</p>
<p>Does anyone have any tips as to how I could go about debugging my code, and where the nan error could be coming from? Thank you!</p>