Add type to VGAE.embed.

lotzma · Nov 6, 2024 · 5388117 · 5388117
1 parent 4678681
commit 5388117
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/l2gv2/embedding/embeddings.py b/l2gv2/embedding/embeddings.py
@@ -1,7 +1,7 @@
 """ Module for embedding patches using the VGAE model """
 import torch
 import torch_geometric as tg
-from l2gv2.models import speye, VGAEconv
+from l2gv2.models import speye, VGAEconv, Patch
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -56,7 +56,7 @@ def embed(
         dim: int=100,
         hidden_dim: int=32,
         decoder=None
-    ):
+    ) -> tuple[list[Patch], list[torch.nn.Module]]:
 
         """ TODO: docstring for `embed` 
         Args: