fix mypy errors

HumanCompatibleAI · Sep 20, 2023 · 1d5739e · 1d5739e
1 parent 7da3e8f
commit 1d5739e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py
@@ -2,7 +2,7 @@
 import abc
 import dataclasses
 import logging
-from typing import Iterable, Iterator, Mapping, Optional, Type, overload
+from typing import Iterable, Iterator, Mapping, Optional, Type, List, overload
 
 import numpy as np
 import torch as th
@@ -408,7 +408,10 @@ def train_gen(
         if learn_kwargs is None:
             learn_kwargs = {}
 
-        callbacks = [self.gen_callback]
+        callbacks: List[BaseCallback] = []
+
+        if self.gen_callback:
+            callbacks.append(self.gen_callback)
 
         if isinstance(callback, list):
             callbacks.extend(callback)

diff --git a/src/imitation/scripts/train_adversarial.py b/src/imitation/scripts/train_adversarial.py
@@ -33,7 +33,7 @@ def __init__(
         interval: int,
     ):
         """Creates new Checkpoint callback."""
-        super().__init__(self)
+        super().__init__()
         self.trainer = trainer
         self.log_dir = log_dir
         self.interval = interval