[python]Colab crashes for unknown reason while training gan(pytorch)
I tried to train gan on some monkey pics but it crashes colab for unknown reason if try to train it.
I am using 1370 128*128 monkey images.
I have no idea where the issue might be, please respond
btw the runtime is gpu, so the problem doesn't linked to that
from torch import optim
import torchvision
from torchvision import transforms
import torch, torch.nn as nn

batch_size = 4

generic_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.ToPILImage(),
    transforms.Resize((128,128)),
    transforms.ToTensor(),
    transforms.Normalize((0., 0., 0.), (6, 6, 6)),
    transforms.Grayscale(),
])

trainset=torchvision.datasets.ImageFolder(root='drive/My Drive/monkeys', transform=generic_transform)

trainloader = torch.utils.data.DataLoader(trainset,batch_size=batch_size, shuffle=True)

def _init_weights(m):
    if isinstance(m, nn.Conv2d):
        nn.init.normal_(m.weight, 0.0, 0.02)

def gen_noise(noise_shape, n_samples, device='cuda:0'):
    return torch.randn(noise_shape, n_samples).to(device)

class Discriminator(nn.Module):
 #convolutional discriminator
  def __init__(self) -&gt; None:
    super(Discriminator, self).__init__()
    
    self.hidden_dim = 64

    self.relu = nn.ReLU(inplace=False)
    self.sigmoid = nn.Sigmoid()

    self.conv_1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=32, stride = 2)
    self.maxPooling_1 = nn.MaxPool2d(kernel_size=3)
    self.conv_2 = nn.Conv2d(in_channels=32, out_channels=16, kernel_size=8, stride = 2)
    self.maxPooling_2 = nn.MaxPool2d(kernel_size=2)
    self.linear_layer = nn.Linear(in_features=self.hidden_dim, out_features=1)

  def forward(self, x) -&gt; float:
    self.x = x
    self.x = self.relu(self.conv_1(self.x))
    self.x = self.maxPooling_1(self.x)
    self.x = self.relu(self.conv_2(self.x))
    self.x = self.maxPooling_2(self.x)
    print(self.x.shape)
    self.x = self.x.view(self.x.shape[0],
                         self.x.shape[1]*self.x.shape[2]*self.x.shape[3])
    self.x = self.sigmoid(self.linear_layer(self.x))

    return self.x



 class Generator(nn.Module):
    #fully connected generator
    def __init__(self, hidden_dim, output_dim, z_dim) -&gt; None:
    super(Generator, self).__init__()
    
    self.relu = nn.ReLU(inplace=False)
    self.hidden_dim = hidden_dim
    self.output_dim = output_dim
    self.z_dim = z_dim

    self.linear_layer_1 = nn.Linear(in_features=self.z_dim, out_features=self.hidden_dim)
    self.linear_layer_2 = nn.Linear(in_features=self.hidden_dim, out_features=self.hidden_dim*2)
    self.linear_layer_3 = nn.Linear(in_features=self.hidden_dim*2, out_features=self.output_dim)

    def forward(self, x) -&gt; torch.tensor:
      self.x = x
      self.x = self.relu(self.linear_layer_1(self.x))
      self.x = self.relu(self.linear_layer_2(self.x))
      self.x = self.relu(self.linear_layer_3(self.x))
      return self.x

class GAN():
  def __init__(self, hidden_dim, output_dim, z_dim, criterion, device=&quot;cuda:0&quot;) -&gt; None:
    if device == &quot;cuda:0&quot;:
      assert torch.cuda.is_available(), &quot;apply gpu&quot;

    self.hidden_dim = hidden_dim
    self.output_dim = output_dim
    self.device = device
    self.criterion = criterion
    self.z_dim = torch.tensor(z_dim).long()
    
    self.discriminator = Discriminator().to(self.device)
    self.d_opt = optim.Adam(self.discriminator.parameters(), lr=0.0001)

    self.generator = Generator(hidden_dim=self.hidden_dim, output_dim=self.output_dim, z_dim=self.z_dim).to(self.device)
    self.g_opt = optim.Adam(self.generator.parameters(), lr=0.0001)

    self.generator = self.generator.apply(_init_weights)
    self.discriminator = self.discriminator.apply(_init_weights)

class GAN_Trainer():
  def __init__(self, z_dim, model, device=&quot;cuda:0&quot;) -&gt; None:
    self.device = device
    self.gan = model
    self.z_dim = z_dim

    self._d_mean_train_loss = None
    self._g_mean_train_loss = None
    
  def train(self, batch) -&gt; None:
    print(1)
    self.batch = batch.to(self.device)
          
    self.noise = gen_noise(self.batch.shape[0],self.z_dim).to(self.device)
    self.gan.g_opt.zero_grad()
    self._g_output = self.gan.generator.forward(self.noise.to(self.device))

    self._g_output = self._g_output.view(self.batch.shape[0], 
                                        1, 
                                        torch.sqrt(torch.tensor(self.gan.output_dim)).int(), 
                                        torch.sqrt(torch.tensor(self.gan.output_dim)).int())
    
    print(self._g_output.shape)

    self._d_for_g_pred = self.gan.discriminator.forward(self._g_output) 
    self._g_loss = self.gan.criterion(self._d_for_g_pred, torch.zeros_like(self._d_for_g_pred))
    self._g_loss.backward()
    self.gan.g_opt.step()

    self.gan.d_opt.zero_grad()
    self._d_fake_pred = self.gan.discriminator.forward(self._g_output)
    self._d_fake_loss = self.gan.criterion(self._g_output, torch.zeros_like(self._g_output))

    self._d_real_pred = self.gan.discriminator.forward(self.batch)
    self._d_real_loss = self.gan.criterion(self.batch, torch.ones_like(self.batch))

    self._d_mean_loss = torch.mean(torch.cat((self._d_fake_loss, self._d_real_loss),0))
    self._d_mean_loss.backward(retain_graph=True)
    self.gan.d_opt.step()
    
    self._d_mean_train_loss = self._d_mean_train_loss + self._d_mean_loss.detach()
    self._g_mean_train_loss = self._g_mean_train_loss + self._g_loss.detach()

torch.cuda.empty_cache()

gan = GAN(hidden_dim=1200,
          output_dim=16384, 
          z_dim = 1000, 
          criterion=nn.BCEWithLogitsLoss())

trainer = GAN_Trainer(model=gan, z_dim=1000)

#here is where it crashes
from tqdm import trange
torch.cuda.empty_cache()

image = trainset[0][0].to(&quot;cuda:0&quot;).view(1, 
                                         trainset[0][0].shape[0],
                                         trainset[0][0].shape[1],
                                         trainset[0][0].shape[2])
trainer.train(batch=image)

please help! I am starting to lose my sanity already, thanks! &lt;3

I've debugged your code a bit, and found that the crash is happening at line:
self._d_mean_loss = torch.mean(torch.cat((self._d_fake_loss, self._d_real_loss),0))

I have tried to find out the reason it is crashing, it looks like your operations are not properly done and some inplace operations are changing the graph and causing pytorch to malfunction.
There are some major issues in the code that you need to change in terms of GAN logic. Here:
self._d_real_loss = self.gan.criterion(self.batch, torch.ones_like(self.batch))

Real loss should be based on the output of discriminator self._d_real_pred. In general you should have three backward calls:

You feed the discriminator with real batch and expect it to output 'real' class.
Generate fake images with generator. Feed it to discriminator and expect it to output 'fake' class.
Lastly for generator, you feed the discriminator fake images again and expect it to fail, optimize for 'fake' class so that the generator can learn outputting better looking fakes.

I would highly recommend this tutorial: https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html . It is a great tutorial for understanding how GANs are trained.
Lastly, I would remove all self.... storing in forward and train functions. Any operations are done tensors update the graph by storing tensors as members and modifying them may cause problems and cause gradients to wipe etc.
