Method forward

axlearn/common/trainer_test.py:1124–1130

(self, input_batch: Nested[Tensor])

                    learning_rate=0.1, decouple_weight_decay=True
                ),
            ),
            max_step=8,
            checkpointer=Checkpointer.default_config().set(
                save_policy=config_for_function(every_n_steps_and_last_policy).set(
                    n=3,
                    max_step=8,
                ),
            ),
        )