Prepare for "Fix type-safety of torch.nn.Module instances": wave 2

Summary: See D52890934

Reviewed By: malfet, r-barnes

Differential Revision: D66245100

fbshipit-source-id: 019058106ac7eaacf29c1c55912922ea55894d23
This commit is contained in:
Edward Yang
2024-11-21 11:08:51 -08:00
committed by Facebook GitHub Bot
parent e20cbe9b0e
commit f6c2ca6bfc
23 changed files with 147 additions and 1 deletions

View File

@@ -123,6 +123,7 @@ class ImplicitronOptimizerFactory(OptimizerFactoryBase):
"""
# Get the parameters to optimize
if hasattr(model, "_get_param_groups"): # use the model function
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
p_groups = model._get_param_groups(self.lr, wd=self.weight_decay)
else:
p_groups = [

View File

@@ -395,6 +395,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
):
prefix = f"e{stats.epoch}_it{stats.it[trainmode]}"
if hasattr(model, "visualize"):
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
model.visualize(
viz,
visdom_env_imgs,