refactored model tests

This commit is contained in:
William Falcon 2019-07-24 13:57:34 -04:00
parent ecb68b52f8
commit c26d200c41
2 changed files with 1 additions and 5 deletions

View File

@ -54,10 +54,6 @@ class LightningTemplateModel(LightningModule):
:param x:
:return:
"""
print('-'*100)
print('x: ', x.device)
print('model: ', self.c_d1.weight.device)
print('-'*100)
x = self.c_d1(x)
x = torch.tanh(x)

View File

@ -107,7 +107,7 @@ def main():
trainer = Trainer(
experiment=exp,
checkpoint_callback=checkpoint,
progress_bar=False,
progress_bar=True,
max_nb_epochs=1,
gpus=[0, 1],
distributed_backend='dp',