Error when training: "closure_loss" is NoneType
Hi all,
I am trying to train my LightningModule but I seem to keep getting the error TypeError: unsupported operand type(s) for /: 'NoneType' and 'int' on the line closure_loss = closure_loss / self.accumulate_grad_batches, in the function training_step() in the file training_loop.py.
I think it might be something to do with how I format my LightningModule, so here is what my LightningModule looks like
class HPAModelV1(pl.LightningModule):
  def __init__(self):
    super().__init__()

    #self.lossfunc = F.cross_entropy
    self.lossfunc = F.nll_loss

    self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=3, padding=7)
    self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.conv3 = nn.Conv2d(16, 16, kernel_size=5, stride=1, padding=1)
    self.dense = nn.Linear(16, 19)

  def forward(self, x): #input size is (256, 3, 256, 256)

    x = x.float()
    
    out = self.conv1(x)
    out = F.relu(out)
    out = F.max_pool2d(out, 3) # output is (bs, 16, 30, 30)
    
    out = self.conv2(out)
    out = F.relu(out)
    out = F.max_pool2d(out, 3) # output is (bs, 16, 10, 10)

    out = self.conv3(out)
    out = F.relu(out)
    out = F.max_pool2d(out, 8) # output is (bs, 16, 1, 1)

    # dense layer
    out = out.reshape(out.size()[0], 16)
    out = self.dense(out)

    return out

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


  def training_step(self, batch, batchidx):
    # set labels and data
    x = batch[0]
    y = batch[1]
    

    # compute loop
    preds = self(x)
  
  
    probs = F.softmax(preds, dim=1)
 

    # compute the loss function
    J = self.lossfunc(probs, y)
 
   

    # compute accuracy 
    acc = accuracy(probs, y)

    
    #log for weights and biases
    self.log('training loss (step)', J)
    self.log('training accuracy (step)', acc)
    self.log('mean training loss (epoch)', J, on_step=False, on_epoch=True)
    self.log('mean training accuracy (epoch)', acc, on_step=False, on_epoch=True)



    # add information to the progress bar
    pbar =  {'train_acc': acc, 'train_loss' : J}

    return J, acc

  def validation_step(self, valbatch, valbatchidx):
    # use the same training step on the val set

    valJ, valAcc = self.training_step(valbatch, valbatchidx)

    # log for wb
    self.log('validation loss (step)', valJ)
    self.log('validation accuracy (step)', valAcc)
    self.log('mean validation loss (epoch)', valJ, on_step=False, on_epoch=True)
    self.log('mean validation accuracy (epoch)', valAcc, on_step=False, on_epoch=True)

    return valJ, valAcc

  def validation_epoch_end(self, valStepOutputs):
    pass
And if it may help in diagnosing the cause of the issue, here is the stack trace and output of of the Trainer:
GPU available: False, used: False
TPU available: True, using: 1 TPU cores
Global seed set to 0

  | Name  | Type   | Params
---------------------------------
0 | conv1 | Conv2d | 448   
1 | conv2 | Conv2d | 2.3 K 
2 | conv3 | Conv2d | 6.4 K 
3 | dense | Linear | 323   
---------------------------------
9.5 K     Trainable params
0         Non-trainable params
9.5 K     Total params
0.038     Total estimated model params size (MB)
Epoch 0: 0%
0/7759 [00:02<?, ?it/s]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-29-caf15077ca9b> in <module>()
      2 os.environ['WANDB_CONSOLE'] = 'on'
      3 trainer = Trainer(logger=wbLogger, tpu_cores=1, deterministic=True, max_epochs=epochNum, replace_sampler_ddp=False, num_sanity_val_steps=0)
----> 4 trainer.fit(HPAModelV1(), trainDL, valDL)
      5 
      6 print(time.time() - t0)

23 frames
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    497 
    498         # dispath `start_training` or `start_testing` or `start_predicting`
--> 499         self.dispatch()
    500 
    501         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in dispatch(self)
    544 
    545         else:
--> 546             self.accelerator.start_training(self)
    547 
    548     def train_or_test_or_predict(self):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/accelerators/accelerator.py in start_training(self, trainer)
     71 
     72     def start_training(self, trainer):
---> 73         self.training_type_plugin.start_training(trainer)
     74 
     75     def start_testing(self, trainer):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/tpu_spawn.py in start_training(self, trainer)
    264             del os.environ["XLA_USE_BF16"]
    265         self._close_logger(trainer)
--> 266         xmp.spawn(self.new_process, **self.xmp_spawn_kwargs)
    267 
    268     def start_testing(self, trainer) -> None:

/usr/local/lib/python3.7/dist-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
    384   pf_cfg = _pre_fork_setup(nprocs)
    385   if pf_cfg.num_devices == 1:
--> 386     _start_fn(0, pf_cfg, fn, args)
    387   else:
    388     return torch.multiprocessing.start_processes(

/usr/local/lib/python3.7/dist-packages/torch_xla/distributed/xla_multiprocessing.py in _start_fn(index, pf_cfg, fn, args)
    321   # environment must be fully setup before doing so.
    322   _setup_replication()
--> 323   fn(gindex, *args)
    324 
    325 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/tpu_spawn.py in new_process(self, process_idx, trainer, mp_queue)
     98         self.barrier("pre-run-stage")
     99 
--> 100         results = trainer.train_or_test_or_predict()
    101 
    102         self.__save_end_of_training_weights(self.lightning_module)

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in train_or_test_or_predict(self)
    554 
    555         else:
--> 556             results = self.run_train()
    557 
    558         return results

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
    635                 with self.profiler.profile("run_training_epoch"):
    636                     # run train epoch
--> 637                     self.train_loop.run_training_epoch()
    638 
    639                 if self.max_steps and self.max_steps <= self.global_step:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
    495             # ------------------------------------
    496             with self.trainer.profiler.profile("run_training_batch"):
--> 497                 batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
    498 
    499             # when returning -1 from train_step, we end epoch early

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in run_training_batch(self, batch, batch_idx, dataloader_idx)
    657 
    658                         # optimizer step
--> 659                         self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
    660 
    661                     else:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
    436             on_tpu=self.trainer._device_type == DeviceType.TPU and _TPU_AVAILABLE,
    437             using_native_amp=using_native_amp,
--> 438             using_lbfgs=is_lbfgs,
    439         )
    440 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/lightning.py in optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_native_amp, using_lbfgs)
   1388             # wraps into LightingOptimizer only for running step
   1389             optimizer = LightningOptimizer._to_lightning_optimizer(optimizer, self.trainer, optimizer_idx)
-> 1390         optimizer.step(closure=optimizer_closure)
   1391 
   1392     def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimizer: Optimizer, optimizer_idx: int):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/optimizer.py in step(self, closure, *args, **kwargs)
    212             profiler_name = f"optimizer_step_and_closure_{self._optimizer_idx}"
    213 
--> 214         self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
    215         self._total_optimizer_step_calls += 1
    216 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/optimizer.py in __optimizer_step(self, closure, profiler_name, **kwargs)
    132 
    133         with trainer.profiler.profile(profiler_name):
--> 134             trainer.accelerator.optimizer_step(optimizer, self._optimizer_idx, lambda_closure=closure, **kwargs)
    135 
    136     def step(self, *args, closure: Optional[Callable] = None, **kwargs):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/accelerators/accelerator.py in optimizer_step(self, optimizer, opt_idx, lambda_closure, **kwargs)
    275         )
    276         if make_optimizer_step:
--> 277             self.run_optimizer_step(optimizer, opt_idx, lambda_closure, **kwargs)
    278         self.precision_plugin.post_optimizer_step(optimizer, opt_idx)
    279         self.training_type_plugin.post_optimizer_step(optimizer, opt_idx, **kwargs)

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/accelerators/tpu.py in run_optimizer_step(self, optimizer, optimizer_idx, lambda_closure, **kwargs)
     32 
     33     def run_optimizer_step(self, optimizer: Optimizer, optimizer_idx: int, lambda_closure: Callable, **kwargs):
---> 34         xm.optimizer_step(optimizer, barrier=False, optimizer_args={'closure': lambda_closure, **kwargs})
     35 
     36     def all_gather(self, tensor: Union[torch.Tensor], group: Optional[Any] = None, sync_grads: bool = False):

/usr/local/lib/python3.7/dist-packages/torch_xla/core/xla_model.py in optimizer_step(optimizer, barrier, optimizer_args, groups)
    779   """
    780   reduce_gradients(optimizer, groups=groups)
--> 781   loss = optimizer.step(**optimizer_args)
    782   if barrier:
    783     mark_step()

/usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py in wrapper(*args, **kwargs)
     86                 profile_name = "Optimizer.step#{}.step".format(obj.__class__.__name__)
     87                 with torch.autograd.profiler.record_function(profile_name):
---> 88                     return func(*args, **kwargs)
     89             return wrapper
     90 

/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.__class__():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

/usr/local/lib/python3.7/dist-packages/torch/optim/adam.py in step(self, closure)
     64         if closure is not None:
     65             with torch.enable_grad():
---> 66                 loss = closure()
     67 
     68         for group in self.param_groups:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in train_step_and_backward_closure()
    652                         def train_step_and_backward_closure():
    653                             result = self.training_step_and_backward(
--> 654                                 split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens
    655                             )
    656                             return None if result is None else result.loss

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
    745         with self.trainer.profiler.profile("training_step_and_backward"):
    746             # lightning module hook
--> 747             result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
    748             self._curr_step_result = result
    749 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in training_step(self, split_batch, batch_idx, opt_idx, hiddens)
    325 
    326 
--> 327             closure_loss = closure_loss / self.trainer.accumulate_grad_batches
    328 
    329             # the loss will get scaled for amp. avoid any modifications to it

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

Thank you, and sorry for all the text
A
Hi @adamDhalla, training_step needs to return one of:

Tensor - The loss tensor
dict - A dictionary. Can include any keys, but must include the key 'loss'
None - Training will skip to the next batch

https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#training-step