TypeError in closure_loss = closure_loss / self.accumulate_grad_batches for Cross_entropy loss
🐛 Bug
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py in optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
   1055         # (if accumulate_grad_batches = 1 no effect)
   1056         closure_loss = training_step_output.minimize if is_result_obj else training_step_output.batch_loss
-> 1057         closure_loss = closure_loss / self.accumulate_grad_batches
   1058 
   1059         # the loss will get scaled for amp. avoid any modifications to it

TypeError: unsupported operand type(s) for /: 'NoneType' and 'int'

Code sample
class CustomModel(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.cnn_model = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size = 5),
            nn.ReLU(),
            nn.AvgPool2d(2, stride = 2),
            nn.Conv2d(6, 16, kernel_size = 5),
            nn.ReLU(),
            nn.AvgPool2d(2, stride = 2),
            nn.Conv2d(16,32,kernel_size = 5),
            nn.ReLU(),
            nn.AvgPool2d(2, stride = 2))

        self.fc_model = nn.Sequential(
            nn.Linear(2592, 1024), # (N, 2592) -> (N, 1024)
            nn.ReLU(),
            nn.Linear(1024, 30))  # (N, 1024)  -> (N, 30)) #30 classes

    def forward(self, x):
        x = self.cnn_model(x)
        # print(x.shape) 
        x = x.view(x.size(0), -1)
        # print(x.shape)    
        x = self.fc_model(x)
        # print(x.shape)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        # acc = FM.accuracy(y_hat, y)
        result = pl.TrainResult()
        print('f')
        return result

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        acc = FM.accuracy(y_hat, y)
        result = pl.EvalResult(checkpoint_on=loss)
        result.log('val_loss', loss, prog_bar=True)
        result.log('val_acc', acc, prog_bar=True)
        print('f')
        return result

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        return optimizer

    def train_dataloader(self):
        train_loader = DataLoader(dataset=train_dataset, shuffle=True, batch_size=32)
        # print("Length of the train_loader:", len(train_loader))
        return train_loader

    def val_dataloader(self):
        return DataLoader(dataset=val_dataset, shuffle=False, batch_size=32)

The error occurs when I am fitting the model to train. Using lightning 0.9.0 on colab. I am loading dataset by mounting drive and using torchvision datasets.ImageFolder function.