Context manager for training with a reduced memory ratio compared to the full memory usage. Returns: float: The exact memory ratio used during the forward pass.
(self, mem_ratio=1.0)
| 205 | @abstractmethod |
| 206 | @contextmanager |
| 207 | def with_mem_ratio(self, mem_ratio=1.0) -> float: |
| 208 | """ |
| 209 | Context manager for training with a reduced memory ratio compared to the full memory usage. |
| 210 | |
| 211 | Returns: |
| 212 | float: The exact memory ratio used during the forward pass. |
| 213 | """ |
| 214 | pass |
| 215 | |
| 216 | def register_memory_controller(self, memory_controller: MemoryController): |
| 217 | self._memory_controller = memory_controller |