Skip to content

Commit 1f2584d

Browse files
piyush182004franciscojavierarceo
piyush182004
and
franciscojavierarceo
authored
fix: Pin expected tensor to same device in online retrieval test (#6777)
test_get_online_features computes `device` from torch.cuda.is_available() and passes it when building the expected tensors for `driver_id` and `trips`, but the `avg_orders_day` assertion omits it. OnlineResponse.to_tensor() places its output on cuda when a GPU is present (sdk/python/feast/online_response.py), so on a GPU host the expected tensor stays on cpu and torch.equal raises: RuntimeError: Expected all tensors to be on the same device, but got other is on cpu, different from other tensors on cuda:0 (when checking argument in method wrapper_CUDA__equal) CI runs CPU-only, where `device` resolves to "cpu" and the omission has no effect, so this only reproduces on machines with CUDA available. Verified on a CUDA host: the unmodified test fails with the GPU visible and passes under CUDA_VISIBLE_DEVICES="", while the patched test passes in both cases. Signed-off-by: Piyush Mondal <[email protected]> Co-authored-by: Francisco Javier Arceo <[email protected]>
1 parent 8ab92e8 commit 1f2584d

1 file changed

Lines changed: 3 additions & 1 deletion

File tree

‎sdk/python/tests/unit/online_store/test_online_retrieval.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,9 @@ def test_get_online_features() -> None:
158158

159159
# Feature values
160160
assert tensor_result["lon"] == ["1.0", "1.0"] # String -> not tensor
161-
assert torch.equal(tensor_result["avg_orders_day"], torch.tensor([1.0, 1.0]))
161+
assert torch.equal(
162+
tensor_result["avg_orders_day"], torch.tensor([1.0, 1.0], device=device)
163+
)
162164
assert tensor_result["name"] == ["John", "John"]
163165
assert torch.equal(tensor_result["trips"], torch.tensor([7, 7], device=device))
164166

0 commit comments

Comments
 (0)