[BUG] Cannot run YOLOv5s model using Scheduler class

[oluwatimilehin]

Describe the bug
Hello, I have written a function to run the YOLOv5s model in PyTorchSIM as follows:

def run_yolo(batch, config):
    from Scheduler.scheduler import Scheduler, SchedulerDNNModel, Request

    scheduler = Scheduler(
        num_request_queue=1, engine_select=Scheduler.FIFO_ENGINE, togsim_config=config
    )
    device = scheduler.execution_engine.module.custom_device()

    model = torch.hub.load("ultralytics/yolov5", "yolov5s").model.eval()
    url = "https://ultralytics.com/images/zidane.jpg"

    response = requests.get(url)
    img = Image.open(BytesIO(response.content)).convert("RGB")

    transform = transforms.Compose(
        [
            transforms.Resize((640, 640)),
            transforms.ToTensor(),
        ]
    )

    input = (transform(img).unsqueeze(0)).to(device)  # [1, 3, H, W]
    opt_fn = torch.compile(dynamic=False)(model.to(device))

    SchedulerDNNModel.register_model("yolov5s", opt_fn)
    request = Request("yolov5s", [input], [], request_queue_idx=0)
    scheduler.add_request(request, request_time=0)

    # Run scheduler
    while not scheduler.is_finished():
        with torch.no_grad():
            scheduler.schedule()

    print("Yolo Simulation Done")

This function follows the same structure as the resnet18 example in the experiments folder. However, I get an error in the prepare_model function of the Scheduler class with the trace:

  File "/workspace/PyTorchSim/experiments/yolo.py", line 49, in run_yolo
    scheduler.schedule()
  File "/workspace/PyTorchSim/Scheduler/scheduler.py", line 461, in schedule
    result.append(self.per_schedule(i))
  File "/workspace/PyTorchSim/Scheduler/scheduler.py", line 446, in per_schedule
    self.execution_engine.submit(request_list, request_queue_idx)
  File "/workspace/PyTorchSim/Scheduler/scheduler.py", line 208, in submit
    self.prepare_model(batched_req_model)
  File "/workspace/PyTorchSim/Scheduler/scheduler.py", line 233, in prepare_model
    ret = req_model.model(*input_tensor_list)
  .....
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
TypeError: conv2d() received an invalid combination of arguments - got (generator, Parameter, Parameter, tuple, tuple, tuple, int), but expected one of:
 * (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, tuple of ints padding, tuple of ints dilation, int groups)
      didn't match because some of the arguments have invalid types: (!generator!, !Parameter!, !Parameter!, !tuple of (int, int)!, !tuple of (int, int)!, !tuple of (int, int)!, int)
 * (Tensor input, Tensor weight, Tensor bias, tuple of ints stride, str padding, tuple of ints dilation, int groups)
      didn't match because some of the arguments have invalid types: (!generator!, !Parameter!, !Parameter!, !tuple of (int, int)!, !tuple of (int, int)!, !tuple of (int, int)!, int)

Looks like we're not passing the correct parameters to the layer, and I've had no luck trying to fix this.

If it helps, PyTorchSIM runs the model when I set it up like:

def run_yolo(batch, config):
    device = PyTorchSimRunner.setup_device().custom_device()
    
    model = torch.hub.load("ultralytics/yolov5", "yolov5s")
    url = "https://ultralytics.com/images/zidane.jpg"
    
    response = requests.get(url)
    img = Image.open(BytesIO(response.content)).convert("RGB")
    
    transform = transforms.Compose([
        transforms.Resize((640, 640)),
        transforms.ToTensor(),
    ])
    
    x = transform(img).unsqueeze(0)   # [1, 3, H, W]
    x = x.to(device)
    

    model.to(device)
    x = x.to(device)
    
    # Compile and run the model with PyTorchSim
    compiled_model = torch.compile(dynamic=False)(model)
    y = compiled_model(x)
    print("Yolo Simulation Done")

However, this provides the simulation results in separate log files and does not accumulate results. It appears that using the Scheduler class is the easiest way to get the results in one place, but please correct me if I'm wrong.

To Reproduce
Find the full script here: https://gist.github.com/oluwatimilehin/c7a0bc759f954eaf81d5aa48b63cbeb5. Note that this is a follow-up to #203, so you might need to make a similar change to reproduce this.