Merge branch 'main' of https://github.com/Yu-Zhewen/fpgaconvnet-torch …

…into main
Yu-Zhewen · Jan 10, 2024 · cd0bab9 · cd0bab9
2 parents 51347eb + ce85a7b
commit cd0bab9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/segmentation/brats2020.py b/models/segmentation/brats2020.py
@@ -88,7 +88,7 @@ def onnx_exporter(self, onnx_path):
         replace_dict = {}
         for module in self.model.modules():
             if isinstance(module, nn.GroupNorm):
-                replace_dict[module] = nn.BatchNorm3d(module.num_channels)
+                replace_dict[module] = nn.BatchNorm3d(module.num_channels).to(self.device)
         self.replace_modules(replace_dict)
         torch.onnx.export(self, random_input, onnx_path, verbose=False, keep_initializers_as_inputs=True)
         model = onnx.load(onnx_path)