Add CPU fallback for segment_csr to run on HF Space

models/ansysLPFMs.py

@@ -8,48 +8,7 @@ import einops
 import torch
 from torch import nn
 from torch_geometric.nn.pool import radius_graph
-# from torch_scatter import segment_csr
-# Try to use torch_scatter if available (cluster / GPU env),
-# otherwise fall back to a pure PyTorch CPU implementation.
-try:
-    from torch_scatter import segment_csr as _segment_csr
-    HAS_TORCH_SCATTER = True
-except Exception:
-    HAS_TORCH_SCATTER = False
-
-    def _segment_csr(src: torch.Tensor, indptr: torch.Tensor, reduce: str = "sum") -> torch.Tensor:
-        """
-        CPU-only fallback for torch_scatter.segment_csr.
-
-        Parameters
-        ----------
-        src : Tensor
-            Shape [N, ...], values to be reduced.
-        indptr : Tensor
-            Shape [num_segments + 1], CSR-style pointers.
-        reduce : str
-            Only 'sum' is implemented here.
-
-        Returns
-        -------
-        out : Tensor
-            Shape [num_segments, ...], segment-wise reductions.
-        """
-        if reduce != "sum":
-            raise NotImplementedError("Fallback segment_csr currently supports only reduce='sum'")
-
-        # Number of segments
-        num_segments = indptr.numel() - 1
-        out_shape = (num_segments,) + src.shape[1:]
-        out = src.new_zeros(out_shape)
-
-        # Simple loop over segments (fine for CPU inference in HF Space)
-        for i in range(num_segments):
-            start = int(indptr[i].item())
-            end = int(indptr[i + 1].item())
-            if start < end:
-                out[i] = src[start:end].sum(dim=0)
-        return out
+from torch_scatter import segment_csr
 import torch.nn.functional as F