Skip to content

Commit 8c83b41

Browse files
EddyLXJmeta-codesync[bot]
authored andcommitted
Support eval mode for st publish (#5085)
Summary: Pull Request resolved: #5085 X-link: https://github.com/facebookresearch/FBGEMM/pull/2093 As title, in silvertorch bulk eval, they will not call eval() for the module but using torch.no_grad() to run. https://www.internalfb.com/code/fbsource/[324dbccd0ab0]/fbcode/dper_lib/silvertorch/core/publish/data_processing/bulk_eval_dmp_gpu.py?lines=1057 So set a eval mode to turn the self.training to False in tbe for bulk eval. Reviewed By: emlin Differential Revision: D86220286 fbshipit-source-id: 9a48c7b4dc09767c99a545d1f25e53bf4265079f
1 parent 9df97a7 commit 8c83b41

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4762,3 +4762,9 @@ def is_first_tbe() -> bool:
47624762
logging.info(
47634763
f"[FREE_MEM Eviction] Evict all at batch {self.step}, {free_cpu_mem_gb} GB free CPU memory, {global_evict_trigger} ranks triggered eviction"
47644764
)
4765+
4766+
def reset_inference_mode(self) -> None:
4767+
"""
4768+
Reset the inference mode
4769+
"""
4770+
self.eval()

0 commit comments

Comments
 (0)