fix python3 int/float error

timomernick · Nov 8, 2017 · b193a4c · b193a4c
1 parent 554f8b7
commit b193a4c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/capsule_network.py b/capsule_network.py
@@ -50,9 +50,9 @@ def __init__(self,
                                    use_routing=True)
 
         reconstruction_size = image_width * image_height * image_channels
-        self.reconstruct0 = nn.Linear(num_output_units*output_unit_size, (reconstruction_size * 2) / 3)
-        self.reconstruct1 = nn.Linear((reconstruction_size * 2) / 3, (reconstruction_size * 3) / 2)
-        self.reconstruct2 = nn.Linear((reconstruction_size * 3) / 2, reconstruction_size)
+        self.reconstruct0 = nn.Linear(num_output_units*output_unit_size, int((reconstruction_size * 2) / 3))
+        self.reconstruct1 = nn.Linear(int((reconstruction_size * 2) / 3), int((reconstruction_size * 3) / 2))
+        self.reconstruct2 = nn.Linear(int((reconstruction_size * 3) / 2), reconstruction_size)
 
         self.relu = nn.ReLU(inplace=True)
         self.sigmoid = nn.Sigmoid()