mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-10-29 13:30:57 +08:00
cleaned up demos
This commit is contained in:
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()
|
||||
Reference in new issue
Block a user