Unverified Commit 193fbb0b authored by fzyzcjy's avatar fzyzcjy Committed by GitHub
Browse files

Super tiny add UT for copy_to_gpu_no_ce (#12270)

parent e607850f
......@@ -369,7 +369,7 @@ def downcast_fp8(
)
def copy_to_gpu_no_ce(input: List[int], output: torch.Tensor):
def copy_to_gpu_no_ce(input: torch.Tensor, output: torch.Tensor):
torch.ops.sgl_kernel.copy_to_gpu_no_ce(input, output)
......
import pytest
import sgl_kernel
import torch
from sgl_kernel.elementwise import copy_to_gpu_no_ce
@pytest.mark.parametrize("size", [64, 72])
def test_copy_to_gpu_no_ce(size):
tensor_cpu = torch.randint(0, 1000000, (size,), dtype=torch.int32, device="cpu")
tensor_gpu = torch.empty_like(tensor_cpu, device="cuda")
copy_to_gpu_no_ce(tensor_cpu, tensor_gpu)
assert torch.all(tensor_cpu.cuda() == tensor_gpu)
if __name__ == "__main__":
pytest.main([__file__])
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment