python: minor test improvement (#221)

* python: minor test improvement

Signed-off-by: Tim Stamler <tstamler@nvidia.com>

Author: tstamler

examples/python/blocking_send_recv_example.py

@@ -26,6 +26,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--ip", type=str, required=True)
     parser.add_argument("--port", type=int, default=5555)
+    parser.add_argument("--use_cuda", type=bool, default=False)
     parser.add_argument(
         "--mode",
         type=str,
@@ -47,10 +48,19 @@ def parse_args():
 
     # Allocate memory and register with NIXL
     agent = nixl_agent(args.mode, config)
-    if args.mode == "target":
+    if args.mode == "target" and args.use_cuda:
+        tensors = [
+            torch.ones(10, dtype=torch.float32, device="cuda:0") for _ in range(2)
+        ]
+    elif args.use_cuda:
+        tensors = [
+            torch.zeros(10, dtype=torch.float32, device="cuda:0") for _ in range(2)
+        ]
+    elif args.mode == "target" and not args.use_cuda:
         tensors = [torch.ones(10, dtype=torch.float32) for _ in range(2)]
     else:
         tensors = [torch.zeros(10, dtype=torch.float32) for _ in range(2)]
+
     print(f"{args.mode} Tensors: {tensors}")
 
     reg_descs = agent.register_memory(tensors)

src/api/python/_api.py

@@ -766,7 +766,7 @@ def get_xfer_descs(
             gpu_id = descs.get_device()
             if gpu_id == -1:  # DRAM
                 gpu_id = 0
-            new_descs = nixlBind.nixlRegDList(
+            new_descs = nixlBind.nixlXferDList(
                 self.nixl_mems[mem_type],
                 [(base_addr, region_len, gpu_id)],
                 is_sorted,