MCPcopy Create free account
hub / github.com/OpenDriveLab/ReSim / training_step

Method training_step

sat/vae_modules/autoencoder.py:295–308  ·  view source on GitHub ↗
(self, batch: dict, batch_idx: int)

Source from the content-addressed store, hash-verified

293 raise NotImplementedError(f"Unknown optimizer {optimizer_idx}")
294
295 def training_step(self, batch: dict, batch_idx: int):
296 opts = self.optimizers()
297 if not isinstance(opts, list):
298 # Non-adversarial case
299 opts = [opts]
300 optimizer_idx = batch_idx % len(opts)
301 if self.global_step < self.disc_start_iter:
302 optimizer_idx = 0
303 opt = opts[optimizer_idx]
304 opt.zero_grad()
305 with opt.toggle_model():
306 loss = self.inner_training_step(batch, batch_idx, optimizer_idx=optimizer_idx)
307 self.manual_backward(loss)
308 opt.step()
309
310 def validation_step(self, batch: dict, batch_idx: int) -> Dict:
311 log_dict = self._validation_step(batch, batch_idx)

Callers

nothing calls this directly

Calls 3

inner_training_stepMethod · 0.95
zero_gradMethod · 0.80
stepMethod · 0.45

Tested by

no test coverage detected