def forward(self, x) -> Tensor:
        """First call the backbone, then the model head."""
        # hack to enable backbone to work properly.
        self.model.device = self.device
        return self.model(x)
First call the backbone, then the model head.
src/flash/pointcloud/detection/model.py