Hello everyone, hope you are having a great time. I  tried to implement a simple VAE! but I’m getting this error out of no where!
The code I’m using is this :
class VAE(nn.Module):
    def __init__(self, embedding=100):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 400)
        self.fc1_mu = nn.Linear(400, embedding)
        self.fc1_std = nn.Linear(400, embedding) 

        self.decoder = nn.Sequential( nn.Linear(embedding, 400), 
                                      nn.ReLU(),
                                      nn.Linear(400, 28*28),
                                      nn.Sigmoid())

    def reparamtrization_trick(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = torch.tensor(std.data.new(std.size()).normal_(0,1))
            return eps.mul(std).add(mu)
        else:
            # During the inference, we simply return the mean of the
            # learned distribution for the current input.  We could
            # use a random sample from the distribution, but mu of
            # course has the highest probability.
            return mu

    def forward(self, input):
        output = input.view(input.size(0), -1)
        output = F.relu(self.fc1(output))
        output_mu = self.fc1_mu(output)
        output_std = self.fc1_std(output)
        z = self.reparamtrization_trick(output_mu, output_std)

        # decoder 
        output = self.decoder(z)
        return output, output_mu, output_std



# now lets train :
epochs = 20 
embeddingsize = 100
interval = 1000
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = VAE(embeddingsize).to(device)
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr =0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 20)

for e in range(epochs):
    for i, (imgs, labels) in enumerate(dataloader_train):
        imgs = imgs.to(device)
        preds,mu, logvar = model(imgs)
       # check to see if the values are in [0,1] range: 
        print(f'min: {np.round(preds.min(1)[0].sum().item())} max: {np.round(preds.max(1)[0].sum().item())}')


        # for loss we simply add the reconstruction loss +kl divergance
        loss_recons = criterion(preds, imgs.view(imgs.size(0), -1))
        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # I guess 0.5 is the beta (a multiplier that specifies how large the distribution 
        # should be)
        # - D_{KL} = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        kl = 0.5 * torch.sum(1+ logvar - mu.pow(2) - logvar.exp())
        # Normalise by same number of elements as in reconstruction
        kl/=imgs.size(0) * (28*28)

        loss = loss_recons + kl 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step() 
        if i% interval ==0:
            print(f'epoch ({e}/{epochs}) loss: {loss.item():.6f}'
                  f'KL: {kl.item():.6f} recons {loss_recons.item():.6f}'
                  f'lr: {scheduler.get_lr()}')
    scheduler.step()


and this is the full stacktrace :
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
 in 
    260 interval = 1000
    261 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
--> 262 model = VAE(embeddingsize).to(device)
    263 criterion = nn.BCELoss()
    264 optimizer = torch.optim.Adam(model.parameters(), lr =0.01)

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in to(self, *args, **kwargs)
    384             return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
    385 
--> 386         return self._apply(convert)
    387 
    388     def register_backward_hook(self, hook):

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)
    191     def _apply(self, fn):
    192         for module in self.children():
--> 193             module._apply(fn)
    194 
    195         for param in self._parameters.values():

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in _apply(self, fn)
    197                 # Tensors stored in modules are graph leaves, and we don't
    198                 # want to create copy nodes, so we have to unpack the data.
--> 199                 param.data = fn(param.data)
    200                 if param._grad is not None:
    201                     param._grad.data = fn(param._grad.data)

~\Anaconda3\lib\site-packages\torch\nn\modules\module.py in convert(t)
    382 
    383         def convert(t):
--> 384             return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
    385 
    386         return self._apply(convert)

RuntimeError: CUDA error: device-side assert triggered

As you can see, this is giving me the error upon instantiating a new object from the VAE class!
I have no idea what is happening. Any help is greatly appreciated.
Update:
For no apparent reason, switching to cpu and then cuda again changed the error and now I’m getting :

RuntimeError : Function ‘AddmmBackward’ returned nan values in its 2th output.

This error occurs when using the official VAE example as well. the code and full stack trace is given below