test(sagemaker): update pytorch-cifar10-sagemaker example to use a wo…

…rking combination of stuff versions (#466)
wandb · Sep 20, 2023 · 1c3e1b6 · 1c3e1b6
1 parent 6e9aab7
commit 1c3e1b6
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/examples/pytorch/pytorch-cifar10-sagemaker/train.py b/examples/pytorch/pytorch-cifar10-sagemaker/train.py
@@ -50,13 +50,14 @@
 estimator = PyTorch(entry_point="cifar10.py",
                     source_dir=os.getcwd() + "/source",
                     role=role,
-                    framework_version='1.7.1',
-                    py_version='py3',
+                    framework_version='1.12.1',
+                    py_version='py38',
                     instance_count=1,
                     instance_type='ml.c5.xlarge',
                     hyperparameters={
                         'epochs': args.train_epochs,
-                        'momentum': 0.9
+                        'momentum': 0.9,
+                        'num_round': 10,
                     })
 
 hyperparameter_ranges = {