minor fix in gp2Scale on GPUs

lbl-camera · Sep 14, 2023 · 06dabd6 · 06dabd6
1 parent 94db5ed
commit 06dabd6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/fvgp/gp.py b/fvgp/gp.py
@@ -2366,10 +2366,10 @@ def wendland_anisotropic_gp2Scale_gpu(x1,x2, hps, obj):
     x1_dev = torch.from_numpy(x1).to(cuda_device, dtype = torch.float32)
     x2_dev = torch.from_numpy(x2).to(cuda_device, dtype = torch.float32)
     hps_dev = torch.from_numpy(hps).to(cuda_device, dtype = torch.float32)
-    distance_matrix = _get_distance_matrix_gpu(x1_dev,x2_dev,cuda_device,hps_dev)
+    d = _get_distance_matrix_gpu(x1_dev,x2_dev,cuda_device,hps_dev)
     d[d > 1.] = 1.
     kernel = hps[0] * (1.-d)**8 * (35.*d**3 + 25.*d**2 + 8.*d + 1.)
-    return kernel
+    return kernel.cpu().numpy()
 
 
 def _get_distance_matrix_gpu(x1,x2,device,hps):