Fix decoding bug

jpeg_codec/codec.py

@@ -25,8 +25,8 @@ def encode(batch, block_size=(8, 8), device=None):
 
 
 def decode(batch, device=None):
-    block_size = np.sqrt(batch.shape[4])
-    image_size = (batch.shape[2] * block_size, batch.shape[3] * block_size)
+    block_size = int(np.sqrt(batch.shape[4]))
+    image_size = (int(batch.shape[2] * block_size), int(batch.shape[3] * block_size))
 
     _, J_i = codec(image_size, (block_size, block_size))