[Piecewise] Support PCG weak_ref_tensor cuda kernel on AMD (#17291)

2026-01-20 14:05:32 -08:00
parent 20ed3822bb
commit 38c233fd04
3 changed files with 9 additions and 2 deletions
@@ -231,6 +231,12 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
      "float scale,"
      "int max_period) -> Tensor");
  m.impl("timestep_embedding", torch::kCUDA, &timestep_embedding);
+
+  /*
+   * From csrc/memory
+   */
+  m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
+  m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
 }

 REGISTER_EXTENSION(common_ops)