fixes

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

python/tutorials/01-vector-add.py

@@ -62,8 +62,7 @@ def add_kernel(x_ptr,  # *Pointer* to first input vector.
 def add(x: torch.Tensor, y: torch.Tensor):
     # We need to preallocate the output.
     output = torch.empty_like(x)
-    is_dvc = f"is_{DEVICE}"
-    assert getattr(x, is_dvc) and getattr(y, is_dvc) and getattr(output, is_dvc)
+    assert x.device.type == DEVICE and y.device.type == DEVICE and output.device.type == DEVICE
     n_elements = output.numel()
     # The SPMD launch grid denotes the number of kernel instances that run in parallel.
     # It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].

python/tutorials/02-fused-softmax.py

@@ -112,7 +112,8 @@ def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n
 # %%
 # We can create a helper function that enqueues the kernel and its (meta-)arguments for any given input tensor.
 
-properties = driver.active.utils.get_device_properties(DEVICE)
+device = getattr(torch, DEVICE).current_device()
+properties = driver.active.utils.get_device_properties(device)
 NUM_SM = properties["multiprocessor_count"]
 SIZE_SMEM = properties["max_shared_mem"]
 WARPS_PER_EU = 8  # TODO: Get from properties