Merge branch 'master' into sampler-docs

pyg-team · Sep 29, 2022 · 6c2bc85 · 6c2bc85
2 parents 390e9a7 + 845ebbf
commit 6c2bc85
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
@@ -145,12 +145,16 @@ def initialize_parameters(self, module, input):
         delattr(self, '_hook')
 
     def _save_to_state_dict(self, destination, prefix, keep_vars):
-        if is_uninitialized_parameter(self.weight):
+        if (is_uninitialized_parameter(self.weight)
+                or torch.onnx.is_in_onnx_export()):
             destination[prefix + 'weight'] = self.weight
         else:
             destination[prefix + 'weight'] = self.weight.detach()
         if self.bias is not None:
-            destination[prefix + 'bias'] = self.bias.detach()
+            if torch.onnx.is_in_onnx_export():
+                destination[prefix + 'bias'] = self.bias
+            else:
+                destination[prefix + 'bias'] = self.bias.detach()
 
     def _lazy_load_hook(self, state_dict, prefix, local_metadata, strict,
                         missing_keys, unexpected_keys, error_msgs):