Add synchronize in aug_fwd

JuliaGPU · Feb 14, 2024 · e92c3b8 · e92c3b8
1 parent 24beeac
commit e92c3b8
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl
@@ -236,9 +236,8 @@ module EnzymeExt
         TapeType, subtape, aug_kernel = _create_tape_kernel(
             kernel, ModifiedBetween, FT, ctxTy, ndrange, iterspace, args2...
         )
-        @show TapeType
         aug_kernel(f, ModifiedBetween, subtape, Val(TapeType), args2...; ndrange, workgroupsize)
-        # KernelAbstractions.synchronize(backend(kernel))
+        KernelAbstractions.synchronize(backend(kernel))
 
         # TODO the fact that ctxTy is type unstable means this is all type unstable.
         # Since custom rules require a fixed return type, explicitly cast to Any, rather