run black

aws · Oct 11, 2024 · 827a2f5 · 827a2f5
1 parent 3dd25c4
commit 827a2f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/sagemaker_tests/huggingface_pytorch/training/resources/scripts/train.py b/test/sagemaker_tests/huggingface_pytorch/training/resources/scripts/train.py
@@ -87,11 +87,11 @@ def run_with_timeout(func, timeout=5):
 
     train_dataset = run_with_timeout(
         train_dataset.map(tokenize, num_proc=1, batched=True, batch_size=len(train_dataset)),
-        timeout=10
+        timeout=10,
     )
     test_dataset = run_with_timeout(
         test_dataset.map(tokenize, num_proc=1, batched=True, batch_size=len(test_dataset)),
-        timeout=10
+        timeout=10,
     )
 
     # set format for pytorch