We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 9df97a7 commit 8c83b41Copy full SHA for 8c83b41
fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py
@@ -4762,3 +4762,9 @@ def is_first_tbe() -> bool:
4762
logging.info(
4763
f"[FREE_MEM Eviction] Evict all at batch {self.step}, {free_cpu_mem_gb} GB free CPU memory, {global_evict_trigger} ranks triggered eviction"
4764
)
4765
+
4766
+ def reset_inference_mode(self) -> None:
4767
+ """
4768
+ Reset the inference mode
4769
4770
+ self.eval()
0 commit comments