Convert train_x, train_y and _y_variance to device in batch mode?

[bk958178]

What's the issue?

in BaseGPController, we have:

        if self.batch_size is None:
            self.train_x = self.train_x.to(self.device)
            self.train_y = self.train_y.to(self.device)
            self._y_variance = self._y_variance.to(self.device)

i.e., when not in batch mode, the inputs are converted to the device.

Do we need to convert these to a device when we are in batch mode?