15 lines
366 B
Python
15 lines
366 B
Python
"""Callbacks for Pytorch Lighning Modules"""
|
|
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
|
|
|
|
class StopOnNaN(pl.Callback):
|
|
def __init__(self, param):
|
|
super().__init__()
|
|
self.param = param
|
|
|
|
def on_epoch_end(self, trainer, pl_module, logs={}):
|
|
if torch.isnan(self.param).any():
|
|
raise ValueError("NaN encountered. Stopping.")
|