Fixing TorchScript support in BatchNorm (#5614)

* update

* update

* changelog

Author: rusty1s

.github/workflows/full_testing.yml

@@ -40,6 +40,7 @@ jobs:
 
       - name: Install internal dependencies
         run: |
+          pip install pyg-lib -f https://data.pyg.org/whl/nightly/torch-${{ matrix.torch-version }}+cpu.html
           pip install torch-scatter torch-sparse torch-cluster torch-spline-conv -f https://data.pyg.org/whl/torch-${{ matrix.torch-version }}+cpu.html
 
       - name: Install main package

CHANGELOG.md

@@ -40,7 +40,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fix `RGCN+pyg-lib` for `LongTensor` input ([#5610](https://github.com/pyg-team/pytorch_geometric/pull/5610))
 - Improved type hint support ([#5603](https://github.com/pyg-team/pytorch_geometric/pull/5603))
 - Avoid modifying `mode_kwargs` in `MultiAggregation` ([#5601](https://github.com/pyg-team/pytorch_geometric/pull/5601))
-- Changed `BatchNorm` to allow for batches of size one during training ([#5530](https://github.com/pyg-team/pytorch_geometric/pull/5530))
+- Changed `BatchNorm` to allow for batches of size one during training ([#5530](https://github.com/pyg-team/pytorch_geometric/pull/5530), [#5614](https://github.com/pyg-team/pytorch_geometric/pull/5614))
 - Integrated better temporal sampling support by requiring that local neighborhoods are sorted according to time ([#5516](https://github.com/pyg-team/pytorch_geometric/issues/5516), [#5602](https://github.com/pyg-team/pytorch_geometric/issues/5602))
 - Fixed a bug when applying several scalers with `PNAConv` ([#5514](https://github.com/pyg-team/pytorch_geometric/issues/5514))
 - Allow `.` in `ParameterDict` key names ([#5494](https://github.com/pyg-team/pytorch_geometric/pull/5494))

torch_geometric/nn/norm/batch_norm.py

@@ -56,11 +56,16 @@ def reset_parameters(self):
     def forward(self, x: Tensor) -> Tensor:
         """"""
         if self.allow_single_element and x.size(0) <= 1:
-            training = self.module.training
-            self.module.eval()
-            out = self.module(x)
-            self.module.training = training
-            return out
+            return torch.nn.functional.batch_norm(
+                x,
+                self.module.running_mean,
+                self.module.running_var,
+                self.module.weight,
+                self.module.bias,
+                False,  # bn_training
+                0.0,  # momentum
+                self.module.eps,
+            )
         return self.module(x)
 
     def __repr__(self):