Handling exceptions raised by forward during a training step

[jangop]

A model might raise an exception during its forward, as, for example, torchvision's current GeneralizedRCNN does:

    if degenerate_boxes.any():
        ...
        torch._assert(False,...)

https://github.com/pytorch/vision/blob/0610b13ac4af3717f538454a9c6b1f441cb386f3/torchvision/models/detection/generalized_rcnn.py#L95

While one might consider this problematic in general, I would like to get around it by essentially skipping a batch if an exception occurs:

class OuterModel(pl.LightningModule):
    ...
        self.rcnn = torchvision...maskrcnn...
    ...
    def training_step(self, batch, batch_idx):
        try:
            loss_dict = self.rcnn(batch)  # In training mode, RCNN directly returns several losses
        except AssertionError:
            log.exception(...)
            return ???
    return sum(loss for loss in loss_dict.values())

The problem is that PL expects a loss returned. Is there any sensible way to go about this?

[akihironitta]

Returning None from training_step is supported in some cases. See https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#training-step