Method cuda

gpytorch/utils/nearest_neighbors.py:68–75

(self, device=None)

        self.to(device)

    def cuda(self, device=None):
        super().cuda(device=device)
        if self.nnlib == "faiss":
            from faiss import GpuIndexFlatL2, StandardGpuResources

            self.res = StandardGpuResources()
            self.index = [GpuIndexFlatL2(self.res, self.dim) for _ in range(self.batch_shape.numel())]
        return self

    def cpu(self):
        super().cpu()