Add warning when initialising the patchgan discriminator with batchnorm in a distributed environment (#454)

* Add warning

* Update generative/networks/nets/patchgan_discriminator.py

Co-authored-by: Eric Kerfoot <[email protected]>
Signed-off-by: Mark Graham <[email protected]>

* Adds missing warning

---------

Signed-off-by: Mark Graham <[email protected]>
Co-authored-by: Eric Kerfoot <[email protected]>

Author: marksgraham

generative/networks/nets/patchgan_discriminator.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import warnings
 from collections.abc import Sequence
 
 import torch
@@ -218,6 +219,12 @@ def __init__(
         )
 
         self.apply(self.initialise_weights)
+        if norm.lower() == "batch" and torch.distributed.is_initialized():
+            warnings.warn(
+                "WARNING: Discriminator is using BatchNorm and a distributed training environment has been detected. "
+                "To train with DDP, convert discriminator to SyncBatchNorm using "
+                "torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).)"
+            )
 
     def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
         """