warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

docs/conf.py

@@ -14,14 +14,13 @@ import os
 import sys
 from datetime import date
 
-import warp as wp
-
-sys.path.insert(0, os.path.abspath(".."))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+import warp as wp  # noqa: E402
 
 # -- Project information -----------------------------------------------------
 
 project = "Warp"
-copyright = f"{date.today().year}, NVIDIA"
+copyright = f"2022-{date.today().year}, NVIDIA"
 author = "NVIDIA"
 
 version = wp.__version__

examples/env/env_ant.py

@@ -54,7 +54,7 @@ class AntEnvironment(Environment):
             up_axis="y",
         )
         builder.joint_q[7:] = [0.0, 1.0, 0.0, -1.0, 0.0, -1.0, 0.0, 1.0]
-        builder.joint_q[:7] = [0.0, 0.7, 0.0, *wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)]
+        builder.joint_q[:7] = [0.0, 0.7, 0.0, *wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)]
 
 
 if __name__ == "__main__":

examples/env/env_cartpole.py

@@ -39,7 +39,7 @@ class CartpoleEnvironment(Environment):
         wp.sim.parse_urdf(
             os.path.join(os.path.dirname(__file__), "../assets/cartpole.urdf"),
             builder,
-            xform=wp.transform((0.0, 0.0, 0.0), wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)),
+            xform=wp.transform((0.0, 0.0, 0.0), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
             floating=False,
             armature=0.1,
             stiffness=0.0,

examples/env/env_humanoid.py

@@ -58,7 +58,7 @@ class HumanoidEnvironment(Environment):
             up_axis="y",
         )
 
-        builder.joint_q[:7] = [0.0, 1.7, 0.0, *wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)]
+        builder.joint_q[:7] = [0.0, 1.7, 0.0, *wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)]
 
 
 if __name__ == "__main__":

examples/example_dem.py

@@ -118,6 +118,8 @@ def integrate(
 
 class Example:
     def __init__(self, stage):
+        self.device = wp.get_device()
+
         self.frame_dt = 1.0 / 60
         self.frame_count = 400
 
@@ -150,32 +152,32 @@ class Example:
         self.use_graph = wp.get_device().is_cuda
 
         if self.use_graph:
-            wp.capture_begin()
-
-            for _ in range(self.sim_substeps):
-                with wp.ScopedTimer("forces", active=False):
-                    wp.launch(
-                        kernel=apply_forces,
-                        dim=len(self.x),
-                        inputs=[
-                            self.grid.id,
-                            self.x,
-                            self.v,
-                            self.f,
-                            self.point_radius,
-                            self.k_contact,
-                            self.k_damp,
-                            self.k_friction,
-                            self.k_mu,
-                        ],
-                    )
-                    wp.launch(
-                        kernel=integrate,
-                        dim=len(self.x),
-                        inputs=[self.x, self.v, self.f, (0.0, -9.8, 0.0), self.sim_dt, self.inv_mass],
-                    )
-
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                for _ in range(self.sim_substeps):
+                    with wp.ScopedTimer("forces", active=False):
+                        wp.launch(
+                            kernel=apply_forces,
+                            dim=len(self.x),
+                            inputs=[
+                                self.grid.id,
+                                self.x,
+                                self.v,
+                                self.f,
+                                self.point_radius,
+                                self.k_contact,
+                                self.k_damp,
+                                self.k_friction,
+                                self.k_mu,
+                            ],
+                        )
+                        wp.launch(
+                            kernel=integrate,
+                            dim=len(self.x),
+                            inputs=[self.x, self.v, self.f, (0.0, -9.8, 0.0), self.sim_dt, self.inv_mass],
+                        )
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True):

examples/example_diffray.py

@@ -304,6 +304,8 @@ class Example:
     """
 
     def __init__(self, stage=None, rot_array=[0.0, 0.0, 0.0, 1.0], verbose=False):
+        self.device = wp.get_device()
+
         self.verbose = verbose
         cam_pos = wp.vec3(0.0, 0.75, 7.0)
         cam_rot = wp.quat(0.0, 0.0, 0.0, 1.0)
@@ -416,12 +418,14 @@ class Example:
         self.loss = wp.zeros(1, dtype=float, requires_grad=True)
 
         # capture graph
-        wp.capture_begin()
-        self.tape = wp.Tape()
-        with self.tape:
-            self.compute_loss()
-        self.tape.backward(self.loss)
-        self.graph = wp.capture_end()
+        wp.capture_begin(self.device)
+        try:
+            self.tape = wp.Tape()
+            with self.tape:
+                self.compute_loss()
+            self.tape.backward(self.loss)
+        finally:
+            self.graph = wp.capture_end(self.device)
 
         self.optimizer = SGD(
             [self.render_mesh.rot],
@@ -432,33 +436,36 @@ class Example:
         )
 
     def ray_trace(self, is_live=False):
-        with wp.ScopedDevice("cuda:0"):
-            # raycast
-            wp.launch(
-                kernel=draw_kernel,
-                dim=self.num_rays,
-                inputs=[
-                    self.render_mesh,
-                    self.camera,
-                    self.texture,
-                    self.rays_width,
-                    self.rays_height,
-                    self.rays,
-                    self.lights,
-                    self.render_mode,
-                ],
-            )
-
-            # downsample
-            wp.launch(
-                kernel=downsample_kernel,
-                dim=self.num_pixels,
-                inputs=[self.rays, self.pixels, self.rays_width, pow(2, self.num_samples)],
-            )
+        # raycast
+        wp.launch(
+            kernel=draw_kernel,
+            dim=self.num_rays,
+            inputs=[
+                self.render_mesh,
+                self.camera,
+                self.texture,
+                self.rays_width,
+                self.rays_height,
+                self.rays,
+                self.lights,
+                self.render_mode,
+            ],
+            device=self.device,
+        )
+
+        # downsample
+        wp.launch(
+            kernel=downsample_kernel,
+            dim=self.num_pixels,
+            inputs=[self.rays, self.pixels, self.rays_width, pow(2, self.num_samples)],
+            device=self.device,
+        )
 
     def compute_loss(self):
         self.ray_trace()
-        wp.launch(loss_kernel, dim=self.num_pixels, inputs=[self.pixels, self.target_pixels, self.loss])
+        wp.launch(
+            loss_kernel, dim=self.num_pixels, inputs=[self.pixels, self.target_pixels, self.loss], device=self.device
+        )
 
     def get_image(self):
         return self.pixels.numpy().reshape((self.height, self.width, 3))

examples/example_fluid.py

@@ -165,6 +165,8 @@ def init(rho: wp.array2d(dtype=float), u: wp.array2d(dtype=wp.vec2), radius: int
 
 class Example:
     def __init__(self, **kwargs):
+        self.device = wp.get_device()
+
         self.sim_fps = 60.0
         self.sim_substeps = 2
         self.iterations = 100
@@ -187,9 +189,11 @@ class Example:
 
         # capture pressure solve as a CUDA graph
         if self.device.is_cuda:
-            wp.capture_begin()
-            self.pressure_iterations()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.pressure_iterations()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
         with wp.ScopedTimer("update"):

examples/example_jacobian_ik.py

@@ -87,7 +87,7 @@ class Example:
             builder.add_builder(
                 articulation_builder,
                 xform=wp.transform(
-                    np.array((i * 2.0, 4.0, 0.0)), wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)
+                    wp.vec3(i * 2.0, 4.0, 0.0), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
                 ),
             )
             self.target_origin.append((i * 2.0, 4.0, 0.0))

examples/example_mesh_intersect.py

@@ -83,6 +83,9 @@ def intersect(
 
 class Example:
     def __init__(self, stage):
+        rng = np.random.default_rng()
+
+        self.device = wp.get_device()
         self.query_count = 1024
         self.has_queried = False
 
@@ -102,10 +105,10 @@ class Example:
 
         for i in range(self.query_count):
             # random offset
-            p = (np.random.rand(3) * 0.5 - 0.5) * 5.0
+            p = wp.vec3(rng.random(3) * 0.5 - 0.5) * 5.0
 
             # random orientation
-            axis = wp.normalize(np.random.rand(3) * 0.5 - 0.5)
+            axis = wp.normalize(wp.vec3(rng.random(3) * 0.5 - 0.5))
             angle = float(np.random.rand(1)[0])
 
             q = wp.quat_from_axis_angle(wp.normalize(axis), angle)
@@ -115,8 +118,8 @@ class Example:
         self.array_result = wp.zeros(self.query_count, dtype=int)
         self.array_xforms = wp.array(self.xforms, dtype=wp.transform)
 
-        # force module load (for accurate profiling)
-        wp.force_load()
+        # compile and load the module up front (for accurate profiling)
+        wp.load_module(device=self.device)
 
     def update(self):
         with wp.ScopedTimer("intersect", active=True):
@@ -143,14 +146,14 @@ class Example:
                     os.path.join(os.path.dirname(__file__), self.path_0),
                     pos=wp.vec3(xform.p[0] + offset, xform.p[1], xform.p[2]),
                     rot=xform.q,
-                    scale=(1.0, 1.0, 1.0),
+                    scale=wp.vec3(1.0, 1.0, 1.0),
                 )
                 self.renderer.render_ref(
                     f"mesh_{i}_1",
                     os.path.join(os.path.dirname(__file__), self.path_1),
-                    pos=(offset, 0.0, 0.0),
+                    pos=wp.vec3(offset, 0.0, 0.0),
                     rot=wp.quat_identity(),
-                    scale=(1.0, 1.0, 1.0),
+                    scale=wp.vec3(1.0, 1.0, 1.0),
                 )
 
                 # if pair intersects then draw a small box above the pair

examples/example_nvdb.py

@@ -143,9 +143,9 @@ class Example:
             self.renderer.render_ref(
                 name="collision",
                 path=os.path.join(os.path.dirname(__file__), "assets/rocks.usd"),
-                pos=(0.0, 0.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi),
-                scale=(1.0, 1.0, 1.0),
+                pos=wp.vec3(0.0, 0.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi),
+                scale=wp.vec3(1.0, 1.0, 1.0),
             )
             self.renderer.render_points(name="points", points=self.positions.numpy(), radius=self.sim_margin)
 

examples/example_render_opengl.py

@@ -27,6 +27,8 @@ split_up_tiles = True
 custom_tile_arrangement = False
 # whether to display the pixels in a matplotlib figure
 show_plot = True
+# whether to render depth image to a Warp array
+render_mode = "depth"
 
 renderer = wp.render.OpenGLRenderer(vsync=False)
 instance_ids = []
@@ -51,11 +53,12 @@ if num_tiles > 1:
 
 renderer.render_ground()
 
+channels = 1 if render_mode == "depth" else 3
 if show_plot:
     import matplotlib.pyplot as plt
 
     if split_up_tiles:
-        pixels = wp.zeros((num_tiles, renderer.tile_height, renderer.tile_width, 3), dtype=wp.float32)
+        pixels = wp.zeros((num_tiles, renderer.tile_height, renderer.tile_width, channels), dtype=wp.float32)
         ncols = int(np.ceil(np.sqrt(num_tiles)))
         nrows = int(np.ceil(num_tiles / float(ncols)))
         img_plots = []
@@ -70,17 +73,23 @@ if show_plot:
             sharey=True,
             num=1,
         )
-        tile_temp = np.zeros((renderer.tile_height, renderer.tile_width, 3), dtype=np.float32)
+        tile_temp = np.zeros((renderer.tile_height, renderer.tile_width, channels), dtype=np.float32)
         for dim in range(ncols * nrows):
             ax = axes[dim // ncols, dim % ncols]
             if dim >= num_tiles:
                 ax.axis("off")
                 continue
-            img_plots.append(ax.imshow(tile_temp))
+            if render_mode == "depth":
+                img_plots.append(ax.imshow(tile_temp, vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane))
+            else:
+                img_plots.append(ax.imshow(tile_temp))
     else:
         fig = plt.figure(1)
-        pixels = wp.zeros((renderer.screen_height, renderer.screen_width, 3), dtype=wp.float32)
-        img_plot = plt.imshow(pixels.numpy())
+        pixels = wp.zeros((renderer.screen_height, renderer.screen_width, channels), dtype=wp.float32)
+        if render_mode == "depth":
+            img_plot = plt.imshow(pixels.numpy(), vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane)
+        else:
+            img_plot = plt.imshow(pixels.numpy())
 
     plt.ion()
     plt.show()
@@ -95,14 +104,14 @@ while renderer.is_running():
     renderer.render_cylinder(
         "cylinder",
         [3.2, 1.0, np.sin(time + 0.5)],
-        np.array(wp.quat_from_axis_angle((1.0, 0.0, 0.0), np.sin(time + 0.5))),
+        np.array(wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.sin(time + 0.5))),
         radius=0.5,
         half_height=0.8,
     )
     renderer.render_cone(
         "cone",
         [-1.2, 1.0, 0.0],
-        np.array(wp.quat_from_axis_angle((0.707, 0.707, 0.0), time)),
+        np.array(wp.quat_from_axis_angle(wp.vec3(0.707, 0.707, 0.0), time)),
         radius=0.5,
         half_height=0.8,
     )
@@ -110,15 +119,15 @@ while renderer.is_running():
 
     if show_plot and plt.fignum_exists(1):
         if split_up_tiles:
-            pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, 3)
+            pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, channels)
         else:
-            pixel_shape = (renderer.screen_height, renderer.screen_width, 3)
+            pixel_shape = (renderer.screen_height, renderer.screen_width, channels)
 
         if pixel_shape != pixels.shape:
             # make sure we resize the pixels array to the right dimensions if the user resizes the window
             pixels = wp.zeros(pixel_shape, dtype=wp.float32)
 
-        renderer.get_pixels(pixels, split_up_tiles=split_up_tiles)
+        renderer.get_pixels(pixels, split_up_tiles=split_up_tiles, mode=render_mode)
 
         if split_up_tiles:
             pixels_np = pixels.numpy()

examples/example_sim_cartpole.py

@@ -28,6 +28,8 @@ wp.init()
 
 class Example:
     def __init__(self, stage=None, num_envs=1, enable_rendering=True, print_timers=True):
+        self.device = wp.get_device()
+
         builder = wp.sim.ModelBuilder()
 
         self.num_envs = num_envs
@@ -37,7 +39,7 @@ class Example:
         wp.sim.parse_urdf(
             os.path.join(os.path.dirname(__file__), "assets/cartpole.urdf"),
             articulation_builder,
-            xform=wp.transform(np.zeros(3), wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.5)),
+            xform=wp.transform(wp.vec3(), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)),
             floating=False,
             density=100,
             armature=0.1,
@@ -98,13 +100,15 @@ class Example:
 
         if self.use_graph:
             # create update graph
-            wp.capture_begin()
-            self.update()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
         with wp.ScopedTimer("simulate", active=True, print=self.print_timers):
-            if self.use_graph is False or self.graph is None:
+            if not self.use_graph or self.graph is None:
                 for _ in range(self.sim_substeps):
                     self.state.clear_forces()
                     self.state = self.integrator.simulate(self.model, self.state, self.state, self.sim_dt)

examples/example_sim_cloth.py

@@ -37,19 +37,10 @@ class IntegratorType(Enum):
 
 
 class Example:
-    def __init__(self, stage):
-        self.parser = argparse.ArgumentParser()
-        self.parser.add_argument(
-            "--integrator",
-            help="Type of integrator",
-            type=IntegratorType,
-            choices=list(IntegratorType),
-            default=IntegratorType.EULER,
-        )
-
-        args = self.parser.parse_args()
+    def __init__(self, stage, integrator=IntegratorType.EULER):
+        self.device = wp.get_device()
 
-        self.integrator_type = args.integrator
+        self.integrator_type = integrator
 
         self.sim_width = 64
         self.sim_height = 32
@@ -68,9 +59,9 @@ class Example:
 
         if self.integrator_type == IntegratorType.EULER:
             builder.add_cloth_grid(
-                pos=(0.0, 4.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi * 0.5),
-                vel=(0.0, 0.0, 0.0),
+                pos=wp.vec3(0.0, 4.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.5),
+                vel=wp.vec3(0.0, 0.0, 0.0),
                 dim_x=self.sim_width,
                 dim_y=self.sim_height,
                 cell_x=0.1,
@@ -83,9 +74,9 @@ class Example:
             )
         else:
             builder.add_cloth_grid(
-                pos=(0.0, 4.0, 0.0),
-                rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), math.pi * 0.5),
-                vel=(0.0, 0.0, 0.0),
+                pos=wp.vec3(0.0, 4.0, 0.0),
+                rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.5),
+                vel=wp.vec3(0.0, 0.0, 0.0),
                 dim_x=self.sim_width,
                 dim_y=self.sim_height,
                 cell_x=0.1,
@@ -109,9 +100,9 @@ class Example:
         builder.add_shape_mesh(
             body=-1,
             mesh=mesh,
-            pos=(1.0, 0.0, 1.0),
-            rot=wp.quat_from_axis_angle((0.0, 1.0, 0.0), math.pi * 0.5),
-            scale=(2.0, 2.0, 2.0),
+            pos=wp.vec3(1.0, 0.0, 1.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.5),
+            scale=wp.vec3(2.0, 2.0, 2.0),
             ke=1.0e2,
             kd=1.0e2,
             kf=1.0e1,
@@ -135,9 +126,11 @@ class Example:
 
         if self.sim_use_graph:
             # create update graph
-            wp.capture_begin()
-            self.update()
-            self.graph = wp.capture_end()
+            wp.capture_begin(self.device)
+            try:
+                self.update()
+            finally:
+                self.graph = wp.capture_end(self.device)
 
     def update(self):
         with wp.ScopedTimer("simulate", dict=self.profiler):
@@ -168,9 +161,20 @@ class Example:
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--integrator",
+        help="Type of integrator",
+        type=IntegratorType,
+        choices=list(IntegratorType),
+        default=IntegratorType.EULER,
+    )
+
+    args = parser.parse_args()
+
     stage_path = os.path.join(os.path.dirname(__file__), "outputs/example_sim_cloth.usd")
 
-    example = Example(stage_path)
+    example = Example(stage_path, integrator=args.integrator)
 
     for i in range(example.sim_frames):
         example.update()

examples/example_sim_fk_grad.py

@@ -84,12 +84,12 @@ class Example:
 
             if i == chain_length - 1:
                 # create end effector
-                builder.add_shape_sphere(pos=(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
+                builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
 
             else:
                 # create shape
                 builder.add_shape_box(
-                    pos=(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
+                    pos=wp.vec3(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
                 )
 
         self.device = wp.get_device(device)

examples/example_sim_fk_grad_torch.py

@@ -95,7 +95,7 @@ class Example:
             builder.add_joint_revolute(
                 parent=parent,
                 child=b,
-                axis=(0.0, 0.0, 1.0),
+                axis=wp.vec3(0.0, 0.0, 1.0),
                 parent_xform=parent_joint_xform,
                 child_xform=wp.transform_identity(),
                 limit_lower=-np.deg2rad(60.0),
@@ -108,12 +108,12 @@ class Example:
 
             if i == chain_length - 1:
                 # create end effector
-                builder.add_shape_sphere(pos=(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
+                builder.add_shape_sphere(pos=wp.vec3(0.0, 0.0, 0.0), radius=0.1, density=10.0, body=b)
 
             else:
                 # create shape
                 builder.add_shape_box(
-                    pos=(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
+                    pos=wp.vec3(chain_width * 0.5, 0.0, 0.0), hx=chain_width * 0.5, hy=0.1, hz=0.1, density=10.0, body=b
                 )
 
         # finalize model

examples/example_sim_grad_bounce.py

@@ -45,6 +45,7 @@ def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha
 
 class Example:
     def __init__(self, stage=None, enable_rendering=True, profile=False, adapter=None, verbose=False):
+        self.device = wp.get_device()
         self.verbose = verbose
 
         # seconds
@@ -71,8 +72,8 @@ class Example:
         mu = 0.2
 
         builder = wp.sim.ModelBuilder()
-        builder.add_particle(pos=(-0.5, 1.0, 0.0), vel=(5.0, -5.0, 0.0), mass=1.0)
-        builder.add_shape_box(body=-1, pos=(2.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0, ke=ke, kf=kf, kd=kd, mu=mu)
+        builder.add_particle(pos=wp.vec3(-0.5, 1.0, 0.0), vel=wp.vec3(5.0, -5.0, 0.0), mass=1.0)
+        builder.add_shape_box(body=-1, pos=wp.vec3(2.0, 1.0, 0.0), hx=0.25, hy=1.0, hz=1.0, ke=ke, kf=kf, kd=kd, mu=mu)
 
         self.device = wp.get_device(adapter)
         self.profile = profile
@@ -106,12 +107,14 @@ class Example:
             self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=1.0)
 
         # capture forward/backward passes
-        wp.capture_begin()
-        self.tape = wp.Tape()
-        with self.tape:
-            self.compute_loss()
-        self.tape.backward(self.loss)
-        self.graph = wp.capture_end()
+        wp.capture_begin(self.device)
+        try:
+            self.tape = wp.Tape()
+            with self.tape:
+                self.compute_loss()
+            self.tape.backward(self.loss)
+        finally:
+            self.graph = wp.capture_end(self.device)
 
     def compute_loss(self):
         # run control loop

examples/example_sim_grad_cloth.py

@@ -53,6 +53,7 @@ def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha
 
 class Example:
     def __init__(self, stage, profile=False, adapter=None, verbose=False):
+        self.device = wp.get_device()
         self.verbose = verbose
 
         # seconds
@@ -80,9 +81,9 @@ class Example:
         dim_y = 16
 
         builder.add_cloth_grid(
-            pos=(0.0, 0.0, 0.0),
-            vel=(0.1, 0.1, 0.0),
-            rot=wp.quat_from_axis_angle((1.0, 0.0, 0.0), -math.pi * 0.25),
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            vel=wp.vec3(0.1, 0.1, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.25),
             dim_x=dim_x,
             dim_y=dim_y,
             cell_x=1.0 / dim_x,
@@ -115,12 +116,14 @@ class Example:
         self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=4.0)
 
         # capture forward/backward passes
-        wp.capture_begin()
-        self.tape = wp.Tape()
-        with self.tape:
-            self.compute_loss()
-        self.tape.backward(self.loss)
-        self.graph = wp.capture_end()
+        wp.capture_begin(self.device)
+        try:
+            self.tape = wp.Tape()
+            with self.tape:
+                self.compute_loss()
+            self.tape.backward(self.loss)
+        finally:
+            self.graph = wp.capture_end(self.device)
 
     def compute_loss(self):
         # run control loop