diff --git a/mu_map/models/unet.py b/mu_map/models/unet.py
index e6e0bbc2ea46910c8be210563d91b52f0396e161..2864d665999eb9452680eb9c0f921050166b0316 100644
--- a/mu_map/models/unet.py
+++ b/mu_map/models/unet.py
@@ -1,5 +1,6 @@
 from typing import Optional, List
 
+import torch
 import torch.nn as nn