cleaned up demos

This commit is contained in:
William Falcon committed 2019-10-05 14:13:32 -04:00
1 parent c86524b0cc
commit 9fc01e3fd3
18 files changed
+91 -283

No files matched your search

@@ -0,0 +1,24 @@
from pytorch_lightning import Trainer
from test_tube import Experiment
import os
def main():
# use the cool model from the main README.md
model = CoolModel() # noqa: F821
exp = Experiment(save_dir=os.getcwd())
# train on 4 GPUs across 4 nodes
trainer = Trainer(
experiment=exp,
distributed_backend='ddp',
max_nb_epochs=10,
gpus=4,
nb_gpu_nodes=4
)
trainer.fit(model)
if __name__ == '__main__':
main()