Gnn recipes bug fixes (#1672)

* gnn recipes bug fixes

* minor fixes

Author: mnabian

examples/additive_manufacturing/sintering_physics/inference.py

@@ -32,15 +32,6 @@
         "Mesh Graph Net Datapipe requires the Tensorflow library. Install the "
         + "package at: https://www.tensorflow.org/install"
     )
-physical_devices = tf.config.list_physical_devices("GPU")
-
-try:
-    for device_ in physical_devices:
-        tf.config.experimental.set_memory_growth(device_, True)
-except:
-    # Invalid device or cannot modify virtual devices once initialized.
-    pass
-
 import hydra
 import torch
 from graph_dataset import GraphDataset

examples/additive_manufacturing/sintering_physics/requirements.txt

@@ -8,4 +8,5 @@ pyvista
 vtk
 natsort
 scikit-learn
-tensorboard
+tensorboard
+tensorflow-cpu>=2.15,<3.0

examples/additive_manufacturing/sintering_physics/train.py

@@ -58,14 +58,6 @@
 )
 from physicsnemo.models.vfgn.graph_network_modules import VFGNLearnedSimulator
 
-physical_devices = tf.config.list_physical_devices("GPU")
-try:
-    for device_ in physical_devices:
-        tf.config.experimental.set_memory_growth(device_, True)
-except:
-    # Invalid device or cannot modify virtual devices once initialized.
-    pass
-
 
 def Train(rank_zero_logger, dist, cfg: DictConfig):
     """
@@ -131,18 +123,13 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
     writer = SummaryWriter(log_dir=cfg.data_options.ckpt_path_vfgn)
 
     optimizer = None
+    scaler = None
     # todo : check device
     device = "cpu"
     step = 0
     running_loss = 0.0
     best_loss = 1000.0
 
-    # Native PyTorch automatic mixed precision (AMP) replaces NVIDIA Apex.
-    # GradScaler is a no-op when enabled=False, so it is safe to construct
-    # unconditionally and only activate when fp16 training is requested.
-    use_amp = cfg.general.fp16
-    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
-
     rank_zero_logger.info("Training started...")
 
     for features, targets in tqdm(dataset):
@@ -184,11 +171,10 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
 
         sampled_noise *= noise_mask
 
-        amp_active = (
-            use_amp and isinstance(device, torch.device) and device.type == "cuda"
-        )
+        amp_enabled = cfg.general.fp16 and scaler is not None
         with torch.autocast(
-            device_type="cuda", dtype=torch.float16, enabled=amp_active
+            device_type=device.type if isinstance(device, torch.device) else "cpu",
+            enabled=amp_enabled,
         ):
             pred_target = model(
                 next_positions=targets.to(device),
@@ -219,8 +205,8 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
             model.setMessagePassingDevices(message_passing_devices)
             model = model.to(device)
             optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
-            # Mixed precision is handled via the torch.amp GradScaler / autocast
-            # constructed above; no extra optimizer wrapping is required.
+            if cfg.general.fp16:
+                scaler = torch.amp.GradScaler(device.type)
 
             scheduler = torch.optim.lr_scheduler.ExponentialLR(
                 optimizer, gamma=0.1, verbose=True
@@ -398,7 +384,7 @@ def Train(rank_zero_logger, dist, cfg: DictConfig):
         rank_zero_logger.info(f"loss: {loss}")
         # back propogation
         optimizer.zero_grad()
-        if use_amp:
+        if cfg.general.fp16:
             scaler.scale(loss).backward()
             scaler.step(optimizer)
             scaler.update()

examples/cfd/external_aerodynamics/xaeronet/README.md

@@ -75,13 +75,30 @@ dataset, please refer to their [paper](https://arxiv.org/pdf/2408.11969).
 
 ## XAeroNet-S prerequisites
 
-Install the requirements using:
+Install the base requirements:
 
 ```bash
 pip install -r requirements.txt
-pip install pyg-lib -f https://data.pyg.org/whl/torch-2.8.0+cu129.html
 ```
 
+`pyg-lib` and `torch_scatter` ship as compiled CUDA extensions and must be
+installed from PyG's pre-built wheel index that matches your installed
+`torch` and CUDA versions. The two-line snippet below detects both and
+constructs the correct URL:
+
+```bash
+TORCH=$(python -c "import torch; print(torch.__version__.split('+')[0])")
+CUDA=$(python -c "import torch; v=torch.version.cuda; \
+print('cu' + v.replace('.', '') if v else 'cpu')")
+pip install pyg-lib torch_scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
+```
+
+If PyG has not published a wheel for your exact torch+CUDA combination
+yet, browse <https://data.pyg.org/whl/> to find the closest match, or
+build from source with `pip install --no-build-isolation torch_scatter`
+(plain `pip install torch_scatter` fails because pip's build isolation
+hides the installed `torch` from the build environment).
+
 See `pyg-lib` [installation instructions](https://github.com/pyg-team/pyg-lib?tab=readme-ov-file#installation)
 for more details.
 

examples/cfd/external_aerodynamics/xaeronet/requirements.txt

@@ -1,6 +1,8 @@
 trimesh>=4.5.0
 torch_geometric>=2.6.1
-torch_scatter>=2.1.2
 pyvista
 vtk
 wandb
+scikit-learn
+tabulate
+matplotlib