Add HyperParameters back to SageMaker training job (#1377)

kubeflow · May 31, 2019 · d5147b9 · d5147b9
1 parent 14c66f0
commit d5147b9
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/components/aws/sagemaker/common/_utils.py b/components/aws/sagemaker/common/_utils.py
@@ -53,6 +53,11 @@ def create_training_job(client, image, instance_type, instance_count, volume_siz
           "VolumeSizeInGB": volume_size
       },
       "TrainingJobName": job_name,
+      "HyperParameters": {	
+          "k": "10",	
+          "feature_dim": "784",	
+          "mini_batch_size": "500"	
+      },
       "StoppingCondition": {
           "MaxRuntimeInSeconds": 60 * 60
       },
@@ -217,4 +222,4 @@ def print_tranformation_job_result(output_location):
   print("Sample transform result: {}".format(results[0]))
 
 def id_generator(size=4, chars=string.ascii_uppercase + string.digits):
-  return ''.join(random.choice(chars) for _ in range(size))
+  return ''.join(random.choice(chars) for _ in range(size))