Update EDVR train_cfg
@@ -32,8 +32,8 @@ class EDVR(BasicRestorer):
pretrained)
self.with_tsa = generator.get('with_tsa', False)
self.step_counter = 0 # count training steps
- if self.with_tsa:
- if self.train_cfg is None or 'tsa_iter' not in self.train_cfg:
+ if self.train_cfg is not None and self.with_tsa:
+ if 'tsa_iter' not in self.train_cfg:
raise KeyError(
'In TSA mode, train_cfg must contain "tsa_iter".')
# only train TSA module at the beginging if with TSA module
