diff --git a/sm_train.py b/sm_train.py
index ccc2a77..408ed57 100644
--- a/sm_train.py
+++ b/sm_train.py
@@ -25,6 +25,7 @@
                     source_dir='code',
                     role=role,
                     framework_version='1.4.0',
+                    py_version='py3',
                     train_instance_count=2,
                     train_instance_type='ml.c4.xlarge',
                     metric_definitions=metric_definitions,