diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py b/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
index 4c38243e4..f19191875 100644
--- a/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
+++ b/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
@@ -49,7 +49,7 @@ def __init__(self,
       ch *= 2
       size = int(size / 2)
     self.conv = ConvBlock(ch, ch * 2, dropout_rate, use_tanh, use_layer_norm, size)
-    size = size/2
+    size = int( size / 2)
 
     self.up_conv = nn.ModuleList()
     self.up_transpose_conv = nn.ModuleList()