Fix minuscule typing error

Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Nov 15, 2024 · d4576d6 · d4576d6
1 parent 1eac9ab
commit d4576d6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/project/algorithms/jax_rl_example_test.py b/project/algorithms/jax_rl_example_test.py
@@ -649,7 +649,7 @@ def log(
         self,
         name: str,
         value: Any,
-        module: JaxRLExample,
+        module: lightning.LightningModule,
         trainer: lightning.Trainer | JaxTrainer,
         **kwargs,
     ):