From 1d5739eff3e050ebc6250259dd33dc1babedaf08 Mon Sep 17 00:00:00 2001 From: Simeon Manolov Date: Tue, 19 Sep 2023 13:33:50 +0300 Subject: [PATCH] fix mypy errors --- src/imitation/algorithms/adversarial/common.py | 7 +++++-- src/imitation/scripts/train_adversarial.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/imitation/algorithms/adversarial/common.py b/src/imitation/algorithms/adversarial/common.py index ab01e4665..dd4bc2f14 100644 --- 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 index 9e1da2d5b..90ba777ad 100644 --- 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