Skip to content

Instantly share code, notes, and snippets.

@yuhanz

yuhanz/torch_array_to_numba.py

Created October 30, 2024 08:16
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save yuhanz/c502edd1cb69169177896aa6083a0574 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save yuhanz/c502edd1cb69169177896aa6083a0574 to your computer and use it in GitHub Desktop.
Download ZIP
Put pytorch array as numba.cuda array
Raw
torch_array_to_numba.py
import numba.cuda
import torch
import numpy as np
t = torch.from_numpy(np.array([1,2,3]))
numba.cuda.as_cuda_array(t.to('cuda'))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment