And here is the training script:

` import os
import sys

from torch import Tensor
sys.path.append("/workspace/")

from pytorch_lightning import Trainer
from pytorch_lightning.loggers import TensorBoardLogger

from easyvision.zoo.edsr.model import EDSR
from model import Model
from dataset import TrainingDataset

from clearml import Task, Dataset

def train():
dataset_path = Dataset.get(
dataset_id="bcd566344203462b839a7ba08dd9efa7"
).get_local_copy()

task = Task.current_task()

params = {
    "batch_size": 8,
    "gpus": 1,

    "auto_select_gpu": True
}

params = task.connect(params)

print(params)
    
dataset = SRDataset(
    path=dataset_path
)

model = Model()

detector_sr = PLWrapper({
    "model": model,
    "dataset_train": dataset,
    "dataset_val": dataset,
    "batch_size_train": params.get("batch_size", 4),
    "batch_size_val": params.get("batch_size", 4),
    "num_workers": params.get("batch_size", 4)
})

trainer = Trainer(
    check_val_every_n_epoch=1,
    num_sanity_val_steps=0,
    gpus=int(params.get("gpus", 1)),
    benchmark=True,
    auto_select_gpus=bool(params.get("auto_select_gpu", True)),
    logger=TensorBoardLogger(save_dir="logs")
)
trainer.fit(detector_sr)

if name == "main":
train() `