Merge pull request #42 from daniel-unyi-42/main

Add Dockerfile

Dockerfile

@@ -0,0 +1,27 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+RUN apt-get update -y && apt-get install -y --no-install-recommends \
+    git \
+    wget \
+    tmux \
+    vim \
+    htop \
+    zip \
+    unzip \
+    build-essential \
+    python3 \
+    python3-pip && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN python3 -m pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir debugpy
+
+WORKDIR /workspace
+
+RUN git clone https://github.com/EliHei2/segger_dev.git /workspace/segger_dev && \
+    pip install -e "/workspace/segger_dev[cuda12,rapids12,cupy12,faiss]"
+
+EXPOSE 5678
+
+ENV PYTHONPATH=/workspace/segger_dev/src:$PYTHONPATH

src/segger/data/parquet/pyg_dataset.py

@@ -70,8 +70,9 @@ def get(self, idx: int) -> Data:
             data["tx"].x = data["tx"].x.unsqueeze(1)
         assert data["tx"].x.dim() == 2
         # this is an issue in PyG's RandomLinkSplit, dimensions are not consistent if there is only one edge in the graph
-        if data["tx", "belongs", "bd"].edge_label_index.dim() == 1:
-            data["tx", "belongs", "bd"].edge_label_index = data["tx", "belongs", "bd"].edge_label_index.unsqueeze(1)
-            data["tx", "belongs", "bd"].edge_label = data["tx", "belongs", "bd"].edge_label.unsqueeze(0)
-        assert data["tx", "belongs", "bd"].edge_label_index.dim() == 2
+        if hasattr(data, "edge_label_index"):
+            if data["tx", "belongs", "bd"].edge_label_index.dim() == 1:
+                data["tx", "belongs", "bd"].edge_label_index = data["tx", "belongs", "bd"].edge_label_index.unsqueeze(1)
+                data["tx", "belongs", "bd"].edge_label = data["tx", "belongs", "bd"].edge_label.unsqueeze(0)
+            assert data["tx", "belongs", "bd"].edge_label_index.dim() == 2
         return data