remove casting

mlcommons · Dec 7, 2023 · ad9b843 · ad9b843
1 parent 515f16b
commit ad9b843
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py b/algorithmic_efficiency/workloads/fastmri/fastmri_pytorch/models.py
@@ -41,27 +41,27 @@ def __init__(self,
     self.size = size
     self.down_sample_layers = nn.ModuleList(
         [ConvBlock(in_chans, num_channels, dropout_rate, use_tanh, use_layer_norm, size)])
-    size = int(size / 2)
+    size = size / 2
     ch = num_channels
     for _ in range(num_pool_layers - 1):
       self.down_sample_layers.append(
           ConvBlock(ch, ch * 2, dropout_rate, use_tanh, use_layer_norm, size))
       ch *= 2
-      size = int(size / 2)
+      size = size / 2
     self.conv = ConvBlock(ch, ch * 2, dropout_rate, use_tanh, use_layer_norm, size)
 
     self.up_conv = nn.ModuleList()
     self.up_transpose_conv = nn.ModuleList()
 
     for _ in range(num_pool_layers - 1):
-      size = int(size * 2)
+      size = size * 2
       self.up_transpose_conv.append(
           TransposeConvBlock(ch * 2, ch, use_tanh, use_layer_norm, size))
       self.up_conv.append(
           ConvBlock(ch * 2, ch, dropout_rate, use_tanh, use_layer_norm, size))
       ch //= 2
 
-    size = int(size * 2)
+    size = size * 2
     self.up_transpose_conv.append(
         TransposeConvBlock(ch * 2, ch, use_tanh, use_layer_norm, size))
     self.up_conv.append(