diff --git a/icenet/data/dataset.py b/icenet/data/dataset.py
index 8763697..d47cc10 100644
--- a/icenet/data/dataset.py
+++ b/icenet/data/dataset.py
@@ -18,6 +18,7 @@
 pytorch_available = False
 try:
     from torch.utils.data import Dataset
+    pytorch_available = True
 except ModuleNotFoundError:
     print("PyTorch not found - not required if not using PyTorch")
 except ImportError: