Hi,
I am trying to train an object detection network which I took from Detectron2 and I loaded it in Pytorch. It has to detect apple. The problem is that when I it gets to the forward() function the following error appears in the fucntion:
train_one_epoch(py_model, optimizer, data_loader, device, epoch, args.print_freq )

File “/home/felipevw/MyAppleDetector/utility/engine.py”, line 30, in train_one_epoch
loss_dict = model(images, targets)
File “/home/felipevw/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
TypeError: forward() takes 2 positional arguments but 3 were given

If someone could hel me, I would be deeply grateful. My code is:

import datetime
import os
import time
import torch
import torch.utils.data
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor


from data.apple_dataset import AppleDataset
from utility.engine import train_one_epoch, evaluate
import utility.utils as utils
import utility.transforms as T


from detectron2.utils.logger import setup_logger
setup_logger()



from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.modeling import build_model

import matplotlib.pyplot as plt


def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)



def get_R50_model(num_classes):
    

    # Get the model from detectron2
    model = get_cfg()
    model.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
    model.DATALOADER.NUM_WORKERS = args.workers
    model.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
    model.MODEL.ROI_HEADS.NUM_CLASSES = 1
    
    py_model = build_model(model)
    
    
    return py_model
def get_R101_model(num_classes):
    
    return 1



def main(args):
    print(args)
    device = args.device
    
    
    # Data loading code
    print("Loading data")
    num_classes = 2
    dataset = AppleDataset(os.path.join(args.data_path, 'train'), get_transform(train=True))
    dataset_test = AppleDataset(os.path.join(args.data_path, 'test'), get_transform(train=False))
    
    
    
    
    print("Creating data loaders")
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True,
                                              num_workers=args.workers, collate_fn=utils.collate_fn)

    data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1,
                                                   shuffle=False, num_workers=args.workers,
                                                   collate_fn=utils.collate_fn)
    
    print("Creating model")
    # Create the correct model type
    if args.model == 'R50_fpn':
        py_model = get_R50_model(num_classes)
    else:
        model = get_R101_model(num_classes)
        
    
    
    
    # Move model to the right device
    py_model.to(device)
    print(py_model)
    
    

    
    
    params = [p for p in py_model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

    #  lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
    
    
    if args.resume:
        checkpoint = torch.load(args.resume, map_location='cpu')
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        
        
    print("Start training")
    start_time = time.time()
    for epoch in range(args.epochs):
        train_one_epoch(py_model, optimizer, data_loader, device, epoch, args.print_freq)
        lr_scheduler.step()

        if args.output_dir:
            torch.save(model.state_dict(), os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))

        # evaluate after every epoch
        evaluate(py_model, data_loader_test, device=device)

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))
    
     
>