diff --git a/tests/test_clearml_resume.py b/tests/test_clearml_resume.py index e9d10cbc0b3..ecb4c8dbd10 100644 --- a/tests/test_clearml_resume.py +++ b/tests/test_clearml_resume.py @@ -1,15 +1,17 @@ import os + +import torch +from clearml import Task +from datasets import load_dataset + from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, - TrainingArguments, - Trainer, DataCollatorWithPadding, + Trainer, + TrainingArguments, ) -from datasets import load_dataset -from clearml import Task -import torch -import shutil + # Set environment variables os.environ["CLEARML_PROJECT"] = "Test Project" @@ -17,11 +19,7 @@ os.environ["CLEARML_TASK"] = "Test Task" os.environ["CLEARML_LOG_MODEL"] = "TRUE" # Initialize ClearML task -task = Task.init( - project_name="Test Project", - task_name="Test Task", - reuse_last_task_id=False -) +task = Task.init(project_name="Test Project", task_name="Test Task", reuse_last_task_id=False) # Load model and tokenizer model_name = "bert-base-uncased" @@ -95,5 +93,5 @@ trainer.train() initial_params = {name: param.data.clone() for name, param in model.named_parameters()} trainer.train() for name, param in model.named_parameters(): - if 'weight' in name: - diff = torch.abs(param.data - initial_params[name]).mean().item() \ No newline at end of file + if "weight" in name: + diff = torch.abs(param.data - initial_params[name]).mean().item()