set dummy labels on example_batch

a-r-j · Dec 26, 2023 · a0782c4 · a0782c4
1 parent b15bdf2
commit a0782c4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/proteinworkshop/datasets/utils.py b/proteinworkshop/datasets/utils.py
@@ -10,6 +10,7 @@
 from typing import List, Optional
 
 import biotite.database.rcsb as rcsb
+import torch
 import torch.nn.functional as F
 from graphein.protein.tensor.data import ProteinBatch, get_random_protein
 from tqdm import tqdm
@@ -136,6 +137,8 @@ def create_example_batch(n: int = 4) -> ProteinBatch:
     batch.pos = batch.coords[:, 1, :]
     batch.x = F.one_hot(batch.residue_type, num_classes=23).float()
 
+    batch.graph_y = torch.randint(0, 2, (n, 1))
+
     batch.x_vector_attr = orientations(batch.pos)
     batch.edge_attr = pos_emb(batch.edge_index, 9)
     batch.edge_vector_attr = _normalize(