Skip to content

Commit

Permalink
Merge pull request #744 from niklasmei/actually_fix
Browse files Browse the repository at this point in the history
exchanged GNN with Model in the checks and in the type hint
  • Loading branch information
niklasmei authored Sep 10, 2024
2 parents 1d47d28 + d836b42 commit dc7fa4f
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions src/graphnet/models/standard_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from torch.optim import Adam

from graphnet.models.gnn.gnn import GNN
from graphnet.models import Model
from .easy_model import EasySyntax
from graphnet.models.task import StandardLearnedTask
from graphnet.models.graphs import GraphDefinition
Expand All @@ -25,7 +26,7 @@ def __init__(
self,
graph_definition: GraphDefinition,
tasks: Union[StandardLearnedTask, List[StandardLearnedTask]],
backbone: GNN = None,
backbone: Model = None,
gnn: Optional[GNN] = None,
optimizer_class: Type[torch.optim.Optimizer] = Adam,
optimizer_kwargs: Optional[Dict] = None,
Expand Down Expand Up @@ -60,7 +61,7 @@ def __init__(
)

# Checks
assert isinstance(backbone, GNN)
assert isinstance(backbone, Model)
assert isinstance(graph_definition, GraphDefinition)

# Member variable(s)
Expand Down

0 comments on commit dc7fa4f

Please sign in to comment.