[Fix] Fix T_max may be None in CosineAnnealingParamScheduler (#735)

pull/436/head
Yixiao Fang 2022-11-19 10:39:19 +08:00 committed by GitHub
parent d6f41bcf74
commit 9d5b417f07
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions

View File

@ -639,7 +639,7 @@ class CosineAnnealingParamScheduler(_ParamScheduler):
@classmethod
def build_iter_from_epoch(cls,
*args,
T_max,
T_max=None,
begin=0,
end=INF,
by_epoch=True,
@ -653,7 +653,8 @@ class CosineAnnealingParamScheduler(_ParamScheduler):
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
T_max = T_max * epoch_length
if T_max is not None:
T_max = T_max * epoch_length
begin = int(begin * epoch_length)
if end != INF:
end = int(end * epoch_length)