diff --git a/SpectralNormalizationKeras.py b/SpectralNormalizationKeras.py
index 623d7a3..8324e36 100644
--- a/SpectralNormalizationKeras.py
+++ b/SpectralNormalizationKeras.py
@@ -17,7 +17,7 @@
 class DenseSN(Dense):
     def build(self, input_shape):
         assert len(input_shape) >= 2
-        input_dim = input_shape[-1]
+        input_dim = int(input_shape[-1])
         self.kernel = self.add_weight(shape=(input_dim, self.units),
                                       initializer=self.kernel_initializer,
                                       name='kernel',
@@ -127,7 +127,7 @@ def build(self, input_shape):
         if input_shape[channel_axis] is None:
             raise ValueError('The channel dimension of the inputs '
                              'should be defined. Found `None`.')
-        input_dim = input_shape[channel_axis]
+        input_dim = int(input_shape[channel_axis])
         kernel_shape = self.kernel_size + (input_dim, self.filters)
 
         self.kernel = self.add_weight(shape=kernel_shape,
@@ -278,7 +278,7 @@ def build(self, input_shape):
         if input_shape[channel_axis] is None:
             raise ValueError('The channel dimension of the inputs '
                              'should be defined. Found `None`.')
-        input_dim = input_shape[channel_axis]
+        input_dim = int(input_shape[channel_axis])
         kernel_shape = self.kernel_size + (input_dim, self.filters)
 
         self.kernel = self.add_weight(shape=kernel_shape,
@@ -357,7 +357,7 @@ def build(self, input_shape):
         if input_shape[channel_axis] is None:
             raise ValueError('The channel dimension of the inputs '
                              'should be defined. Found `None`.')
-        input_dim = input_shape[channel_axis]
+        input_dim = int(input_shape[channel_axis])
         kernel_shape = self.kernel_size + (input_dim, self.filters)
 
         self.kernel = self.add_weight(shape=kernel_shape,
@@ -435,7 +435,7 @@ def build(self, input_shape):
         if input_shape[channel_axis] is None:
             raise ValueError('The channel dimension of the inputs '
                              'should be defined. Found `None`.')
-        input_dim = input_shape[channel_axis]
+        input_dim = int(input_shape[channel_axis])
         kernel_shape = self.kernel_size + (input_dim, self.filters)
 
         self.kernel = self.add_weight(shape=kernel_shape,
@@ -569,7 +569,7 @@ def build(self, input_shape):
         if input_shape[channel_axis] is None:
             raise ValueError('The channel dimension of the inputs '
                              'should be defined. Found `None`.')
-        input_dim = input_shape[channel_axis]
+        input_dim = int(input_shape[channel_axis])
         kernel_shape = self.kernel_size + (self.filters, input_dim)
 
         self.kernel = self.add_weight(shape=kernel_shape,