Where to use `to(self.device)` ?
What is your question?
So currently in my lightning module, I create the multibox module inside the training_step, in order to send the priorboxes to the right device, but when I define the multibox loss in the __init__ method, the priors doesn't seem to be in the right device.
Code
for eg. below works:-
class SSD_simple(pl.LightningModule):

    def __init__(self,config : dict):
        super().__init__()
        self.config = config

        self.model = SSD300()

    
    def forward(self, x):
        return self.model(x)
    
    def training_step(self, batch, batch_nb):
        images, bboxes, labels = batch
        locs, confs = self(images)

        priors = PriorBox(self.config)
        priors = priors.create_priors().to(self.device)
        multibox_loss = MultiBoxLoss(num_classes=self.config['num_classes'],priors=priors,cfg=self.config,device=self.device)

        loc_loss, conf_loss = multibox_loss(locs, confs, bboxes, labels)
        loss = conf_loss + loc_loss

        logs = {'train_loss' : loss, 'conf_loss':conf_loss, 'loc_loss' : loc_loss}

        return {'loss' : loss, 'log' : logs}
But the below doesn't work :-
class SSD_simple(pl.LightningModule):

    def __init__(self,config : dict):
        super().__init__()
        self.config = config

        self.model = SSD300()

        priors = PriorBox(self.config)
        priors = priors.create_priors().to(self.device)
        self.multibox_loss = MultiBoxLoss(num_classes=self.config['num_classes'],priors=priors,cfg=self.config,device=self.device)

    
    def forward(self, x):
        return self.model(x)
    
    def training_step(self, batch, batch_nb):
        images, bboxes, labels = batch
        locs, confs = self(images)

        loc_loss, conf_loss = self.multibox_loss(locs, confs, bboxes, labels)
        loss = conf_loss + loc_loss

        logs = {'train_loss' : loss, 'conf_loss':conf_loss, 'loc_loss' : loc_loss}

        return {'loss' : loss, 'log' : logs}
It throws the following error:
RuntimeError: expected device cuda:0 but got device cpu
I want to get the later method to work since it is not computationally expensive as I am not creating priors again and again.
Any help ? I was thinking whether there was a pre_training method ?
What's your environment?

OS: [e.g. Linux]
Packaging [pip]