Examples use GPUs if available.

prototorch/models/callbacks.py

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