Merge pull request #141 from marpaia/marpaia/fix-fix

Minimally fix typing issues
GFNOrg · Oct 17, 2023 · d9ca558 · d9ca558
2 parents 2784c54 + 13737f6
commit d9ca558
Showing 1 changed file with 2 additions and 4 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, Generic[TrainingSampleType]):
+class PFBasedGFlowNet(GFlowNet[TrainingSampleType]):
     r"""Base class for gflownets that explicitly uses $P_F$.
 
     Attributes:
@@ -75,9 +75,7 @@ def sample_trajectories(self, env: Env, n_samples: int) -> Trajectories:
         return trajectories
 
 
-class TrajectoryBasedGFlowNet(
-    PFBasedGFlowNet[Trajectories], Generic[TrainingSampleType]
-):
+class TrajectoryBasedGFlowNet(PFBasedGFlowNet[Trajectories]):
     def get_pfs_and_pbs(
         self,
         trajectories: Trajectories,