Update Documentation

Clean up project

prototorch/models/callbacks.py

@@ -1,14 +0,0 @@
-"""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.")