Merge pull request #139 from saleml/fix_tests

Fix tests
GFNOrg · Oct 5, 2023 · 2784c54 · 2784c54
2 parents 236f98b + 8af2cbb
commit 2784c54
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/gfn/gflownet/base.py b/src/gfn/gflownet/base.py
@@ -55,7 +55,7 @@ def loss(self, env: Env, training_objects):
         """Computes the loss given the training objects."""
 
 
-class PFBasedGFlowNet(GFlowNet):
+class PFBasedGFlowNet(GFlowNet, Generic[TrainingSampleType]):
     r"""Base class for gflownets that explicitly uses $P_F$.
 
     Attributes:
@@ -75,7 +75,9 @@ def sample_trajectories(self, env: Env, n_samples: int) -> Trajectories:
         return trajectories
 
 
-class TrajectoryBasedGFlowNet(PFBasedGFlowNet[Trajectories]):
+class TrajectoryBasedGFlowNet(
+    PFBasedGFlowNet[Trajectories], Generic[TrainingSampleType]
+):
     def get_pfs_and_pbs(
         self,
         trajectories: Trajectories,