15 lines
393 B
Python
15 lines
393 B
Python
import torch
|
|
from torch.utils.cpp_extension import load
|
|
|
|
# Load the compiled CUDA extension
|
|
my_cuda_extension = load(name='my_cuda_kernel', sources=['my_cuda_kernel.cu'])
|
|
|
|
# Define a PyTorch tensor on the GPU
|
|
input_tensor = torch.cuda.FloatTensor([1, 2, 3, 4])
|
|
|
|
# Call the CUDA function
|
|
output_tensor = my_cuda_extension.my_cuda_function(input_tensor)
|
|
|
|
# Print the result
|
|
print(output_tensor)
|