Fix array assignment issue

james77777778 · May 29, 2024 · 5da58ca · 5da58ca
1 parent 3480c5f
commit 5da58ca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kimm/_src/layers/reparameterizable_conv2d.py b/kimm/_src/layers/reparameterizable_conv2d.py
@@ -306,8 +306,8 @@ def _get_reparameterized_weights_from_layer(self, layer):
         elif isinstance(layer, layers.BatchNormalization):
             k = self.kernel_size[0]
             input_filters = 1 if self.use_depthwise else self.input_filters
-            kernel = ops.convert_to_numpy(
-                ops.zeros([k, k, input_filters, self.filters])
+            kernel = np.zeros(
+                shape=[k, k, input_filters, self.filters], dtype="float64"
             )
             for i in range(self.input_filters):
                 group_i = 0 if self.use_depthwise else i