Fixed saving models (#36)

Author: spietras

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/module/module.py

@@ -78,6 +78,8 @@ async def _build_default_state(self) -> State:
     @staticmethod
     async def _save_models(models: ModelsRegistry) -> None:
         await models.policy.save()
+        await models.value.save()
+        await models.baseline.save()
 
     @classmethod
     async def _save_trainer(cls, trainer: Trainer, directory: Path) -> None:
@@ -205,6 +207,7 @@ async def reset_self(self) -> None:
         async with self.state.write_lock() as state:
             await state.models.policy.reset()
             await state.models.value.reset()
+            await state.models.baseline.reset()
         await super().reset_self()
 
     async def get_metrics(self) -> Collection[Metric]: