Skip to content

Commit

Permalink
prevent side effects in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
awaelchli committed Jun 8, 2020
1 parent a1ddb86 commit cc4cb6b
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions tests/core/test_decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,13 @@
def test_auto_move_data(src_device, dest_device):
""" Test that the decorator moves the data to the device the model is on. """

class CurrentModel(EvalModelTemplate):
pass

# apply the decorator
EvalModelTemplate.forward = auto_move_data(EvalModelTemplate.forward)
CurrentModel.forward = auto_move_data(CurrentModel.forward)

model = EvalModelTemplate()
model = CurrentModel()
model = model.to(dest_device)
model.prepare_data()
loader = model.train_dataloader()
Expand Down

0 comments on commit cc4cb6b

Please sign in to comment.