From 9d6cd91763b84234305545f96160851a038e1b99 Mon Sep 17 00:00:00 2001 From: Gal Rotem Date: Tue, 19 Sep 2023 14:14:16 -0700 Subject: [PATCH] add use_buffers to swa_utils interface (#541) Summary: Pull Request resolved: https://github.com/pytorch/tnt/pull/541 X-link: https://github.com/pytorch/pytorch/pull/109078 As title, this already exists in swa_utils.py Reviewed By: daniellepintz Differential Revision: D49155243 fbshipit-source-id: df0e071247fa79d34eefae624bf46b896c62cb10 --- torchtnt/framework/auto_unit.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/torchtnt/framework/auto_unit.py b/torchtnt/framework/auto_unit.py index af529de1d9..e357e2aaff 100644 --- a/torchtnt/framework/auto_unit.py +++ b/torchtnt/framework/auto_unit.py @@ -54,10 +54,8 @@ def __call__(self, *args, **kwargs): if x.swa_params: if not x.swa_params.avg_fn: - # pyre-ignore: Unexpected keyword [28] x.swa_model = AveragedModel(x.module, use_buffers=True) else: - # pyre-ignore: Unexpected keyword [28] x.swa_model = AveragedModel( x.module, avg_fn=x.swa_params.avg_fn, use_buffers=True )