explicit float32 tensors for property encoders

Author: aditya0by0

chebai_graph/preprocessing/datasets/chebi.py

@@ -443,7 +443,7 @@ def _merge_props_into_base(
         assert (
             max_len_node_properties is not None
         ), "Maximum len of node properties should not be None"
-        x = torch.zeros((num_nodes, max_len_node_properties))
+        x = torch.zeros((num_nodes, max_len_node_properties), dtype=torch.float32)
 
         # Track column offsets for each node type
         atom_offset, fg_offset, graph_offset = 0, 0, 0

chebai_graph/preprocessing/property_encoder.py

@@ -54,7 +54,7 @@ def on_finish(self) -> None:
         return
 
 
-class IndexEncoder(PropertyEncoder):
+class IndexEncoder(PropertyEncoder, abc.ABC):
     """
     Encodes property values as indices. For that purpose, compiles a dynamic list of different values that have
     occurred. Stores this list in a file for later reference.
@@ -148,11 +148,11 @@ def encode(self, token: str | None) -> torch.Tensor:
         """
         if token is None:
             self._count_for_unk_token += 1
-            return torch.tensor([self._unk_token_idx])
+            return torch.tensor([self._unk_token_idx], dtype=torch.float32)
 
         if str(token) not in self.cache:
             self.cache[str(token)] = len(self.cache)
-        return torch.tensor([self.cache[str(token)] + self.offset])
+        return torch.tensor([self.cache[str(token)] + self.offset], dtype=torch.float32)
 
 
 class OneHotEncoder(IndexEncoder):
@@ -215,11 +215,11 @@ def encode(self, token: str | None) -> torch.Tensor:
         """
         if token not in self.tokens_dict:
             self._count_for_unk_token += 1
-            return torch.zeros(1, self.get_encoding_length(), dtype=torch.int64)
+            return torch.zeros(1, self.get_encoding_length(), dtype=torch.float32)
 
         return torch.nn.functional.one_hot(
             self.tokens_dict[token], num_classes=self.get_encoding_length()
-        )
+        ).to(dtype=torch.float32)
 
 
 class AsIsEncoder(PropertyEncoder):
@@ -243,8 +243,8 @@ def encode(self, token: float | int | None) -> torch.Tensor:
             Tensor of shape (1,) containing the input value or zero.
         """
         if token is None:
-            return torch.tensor([0])
-        return torch.tensor([token])
+            return torch.tensor([0], dtype=torch.float32)
+        return torch.tensor([token], dtype=torch.float32)
 
 
 class BoolEncoder(PropertyEncoder):
@@ -267,4 +267,4 @@ def encode(self, token: bool) -> torch.Tensor:
         Returns:
             Tensor with 1 if True else 0.
         """
-        return torch.tensor([1 if token else 0])
+        return torch.tensor([1 if token else 0], dtype=torch.float32)