torch20 lib

dockerfiles/Dockerfile_cuda121_torch2.1

@@ -41,14 +41,6 @@ RUN pip install \
     torchvision \
     torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
 
-RUN TORCH_VERSION=`(python -c "import torch;print(torch.__version__)")` &&
-    pip install \
-    torch_geometric \
-    torch_scatter \
-    torch_sparse \
-    torch_cluster \
-    torch_spline_conv --extra-index-url https://data.pyg.org/whl/torch-${TORCH_VERSION}.html
-
 RUN GDAL_VERSION=$(gdal-config --version | awk -F'[.]' '{print $1"."$2"."$3}') && \
     pip install GDAL==$GDAL_VERSION --no-binary=gdal
 

src/cultionet/augment/augmenters.py

@@ -394,7 +394,7 @@ class Augmenters:
                         each labeled parcel in `y`.
 
             aug_args: Additional keyword arguments passed to the
-                `torch_geometric.data.Data` object.
+                `Data` object.
 
     Example:
         >>> augmenters = Augmenters(augmentations=['tswarp'])

src/cultionet/utils/normalize.py

@@ -66,7 +66,7 @@ def transform(self, batch: Data) -> Data:
         r"""Normalizes data by the Dynamic World log method or by z-scores.
 
         Args:
-            batch (Data): A `torch_geometric` data object.
+            batch (Data): A tensor data object.
             data_means (Tensor): The data feature-wise means.
             data_stds (Tensor): The data feature-wise standard deviations.
 

tests/test_train.py

@@ -3,9 +3,9 @@
 import joblib
 import pytorch_lightning as pl
 import torch
-from torch_geometric.data import Data
 
 import cultionet
+from cultionet.data.data import Data
 from cultionet.data.datasets import EdgeDataset
 from cultionet.enums import AttentionTypes, ModelTypes, ResBlockTypes
 from cultionet.model import CultionetParams