Merge branch 'main' into main

datamol-io · Aug 22, 2023 · f3c9711 · f3c9711
2 parents f74bf20 + f96c1b6
commit f3c9711
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.8", "3.9", "3.10"]
         pytorch-version: ["2.0"]
 
     runs-on: "ubuntu-latest"

diff --git a/graphium/finetuning/finetuning_architecture.py b/graphium/finetuning/finetuning_architecture.py
@@ -1,9 +1,7 @@
-from copy import deepcopy
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import torch
 import torch.nn as nn
-from loguru import logger
 from torch import Tensor
 from torch_geometric.data import Batch
 
@@ -307,7 +305,7 @@ def __init__(self, finetuning_head_kwargs: Dict[str, Any]):
         self.net = net(**finetuning_head_kwargs)
 
     def forward(self, g: Union[Dict[str, Union[torch.Tensor, Batch]], torch.Tensor, Batch]):
-        if isinstance(g, Union[torch.Tensor, Batch]):
+        if isinstance(g, (torch.Tensor, Batch)):
             pass
         elif isinstance(g, Dict) and len(g) == 1:
             g = list(g.values())[0]