Add a device sync at the end

src/targets/gpu/hip.cpp

@@ -88,6 +88,11 @@ argument from_gpu(argument arg)
     return result;
 }
 
+void gpu_sync()
+{
+    hipDeviceSynchronize();
+}
+
 } // namespace gpu
 
 } // namespace migraph

src/targets/gpu/include/migraph/gpu/hip.hpp

@@ -13,6 +13,8 @@ migraph::argument to_gpu(migraph::argument arg, bool host = false);
 
 migraph::argument from_gpu(migraph::argument arg);
 
+void gpu_sync();
+
 struct hip_allocate
 {
     std::string tag{};
@@ -44,6 +46,23 @@ struct hip_load
     }
 };
 
+struct hip_sync
+{
+    std::string tag{};
+    std::string name() const { return "hip::sync"; }
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        if(inputs.empty()) return {};
+        else return inputs.front();
+    }
+    argument compute(context&, const shape&, const std::vector<argument>& args) const
+    {
+        gpu_sync();
+        if(args.empty()) return {};
+        else return args.front();
+    }
+};
+
 struct hip_write
 {
     std::string name() const { return "hip::write"; }

src/targets/gpu/lowering.cpp

@@ -330,6 +330,7 @@ struct miopen_apply
                 check_shape(s, apply_batch_norm_inference(it));
             }
         }
+        prog->insert_instruction(prog->end(), hip_sync{}, std::prev(prog->end()));
     }
 
     instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")