mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-10-11 12:50:09 +08:00
Jirka Borovec
1077159834
* refactor default model * drop redundant seeds * drop redundant seeds * refactor models tests * refactor models tests * imports * fix conf * Apply suggestions from code review
30 lines
809 B
Python
30 lines
809 B
Python
import pytest
|
|
|
|
import tests.base.utils as tutils
|
|
from pytorch_lightning import Trainer
|
|
from tests.base import EvalModelTemplate
|
|
|
|
|
|
@pytest.mark.parametrize('max_steps', [1, 2, 3])
|
|
def test_on_before_zero_grad_called(max_steps):
|
|
|
|
class CurrentTestModel(EvalModelTemplate):
|
|
on_before_zero_grad_called = 0
|
|
|
|
def on_before_zero_grad(self, optimizer):
|
|
self.on_before_zero_grad_called += 1
|
|
|
|
model = CurrentTestModel(tutils.get_default_hparams())
|
|
|
|
trainer = Trainer(
|
|
max_steps=max_steps,
|
|
num_sanity_val_steps=5,
|
|
)
|
|
assert 0 == model.on_before_zero_grad_called
|
|
trainer.fit(model)
|
|
assert max_steps == model.on_before_zero_grad_called
|
|
|
|
model.on_before_zero_grad_called = 0
|
|
trainer.test(model)
|
|
assert 0 == model.on_before_zero_grad_called
|