diff --git a/src/mednet/utils/summary.py b/src/mednet/utils/summary.py
index 2d3824c0fbbcd94399341e6b09118c4f8fe2fa99..ffc3b417033bd927ac91f67919716d68b0ccf02a 100644
--- a/src/mednet/utils/summary.py
+++ b/src/mednet/utils/summary.py
@@ -6,7 +6,7 @@
 
 from functools import reduce
 
-import torch.nn.Module
+import torch
 
 from torch.nn.modules.module import _addindent