diff --git a/model_util/model_api.h b/model_util/model_api.h
index bba7811..22a2300 100644
--- a/model_util/model_api.h
+++ b/model_util/model_api.h
@@ -54,7 +54,8 @@
 iree_hal_executable_library_query_fn_t library_query(void);
 
 // Function to create the bytecode or C module.
-iree_status_t create_module(iree_vm_module_t **module);
+iree_status_t create_module(iree_vm_instance_t *instance,
+                            iree_vm_module_t **module);
 
 // For each ML workload, based on the model configuration, allocate the buffer
 // and prepare the data. It can be loaded from a embedded image binary, a
diff --git a/model_util/util.c b/model_util/util.c
index 2deb726..0e2435a 100644
--- a/model_util/util.c
+++ b/model_util/util.c
@@ -67,25 +67,25 @@
 }
 
 iree_status_t run(const MlModel *model) {
-  IREE_RETURN_IF_ERROR(iree_hal_module_register_all_types());
-
   iree_vm_instance_t *instance = NULL;
   iree_status_t result =
       iree_vm_instance_create(iree_allocator_system(), &instance);
 
+  IREE_RETURN_IF_ERROR(iree_hal_module_register_all_types(instance));
+
   iree_hal_device_t *device = NULL;
   if (iree_status_is_ok(result)) {
     result = create_sample_device(iree_allocator_system(), &device);
   }
   iree_vm_module_t *hal_module = NULL;
   if (iree_status_is_ok(result)) {
-    result = iree_hal_module_create(device, IREE_HAL_MODULE_FLAG_NONE,
+    result = iree_hal_module_create(instance, device, IREE_HAL_MODULE_FLAG_NONE,
                                     iree_allocator_system(), &hal_module);
   }
   // Load bytecode or C module.
   iree_vm_module_t *module = NULL;
   if (iree_status_is_ok(result)) {
-    result = create_module(&module);
+    result = create_module(instance, &module);
   }
 
   // Allocate a context that will hold the module state across invocations.
diff --git a/samples/simple_vec_mul/float_vec.c b/samples/simple_vec_mul/float_vec.c
index a99aa47..90ed92e 100644
--- a/samples/simple_vec_mul/float_vec.c
+++ b/samples/simple_vec_mul/float_vec.c
@@ -43,15 +43,17 @@
     .model_name = "simple_float_vec_mul",
 };
 
-iree_status_t create_module(iree_vm_module_t **module) {
+iree_status_t create_module(iree_vm_instance_t *instance,
+                            iree_vm_module_t **module) {
 #if !defined(BUILD_EMITC)
   const struct iree_file_toc_t *module_file_toc =
       samples_simple_vec_mul_simple_float_mul_bytecode_module_static_create();
   return iree_vm_bytecode_module_create(
+      instance,
       iree_make_const_byte_span(module_file_toc->data, module_file_toc->size),
       iree_allocator_null(), iree_allocator_system(), module);
 #else
-  return module_create(iree_allocator_system(), module);
+  return module_create(instance, iree_allocator_system(), module);
 #endif
 }
 
diff --git a/samples/simple_vec_mul/int_vec.c b/samples/simple_vec_mul/int_vec.c
index 0053c52..63788a8 100644
--- a/samples/simple_vec_mul/int_vec.c
+++ b/samples/simple_vec_mul/int_vec.c
@@ -44,15 +44,17 @@
     .model_name = "simple_int_vec_mul",
 };
 
-iree_status_t create_module(iree_vm_module_t **module) {
+iree_status_t create_module(iree_vm_instance_t *instance,
+                            iree_vm_module_t **module) {
 #if !defined(BUILD_EMITC)
   const struct iree_file_toc_t *module_file_toc =
       samples_simple_vec_mul_simple_int_mul_bytecode_module_static_create();
   return iree_vm_bytecode_module_create(
+      instance,
       iree_make_const_byte_span(module_file_toc->data, module_file_toc->size),
       iree_allocator_null(), iree_allocator_system(), module);
 #else
-  return module_create(iree_allocator_system(), module);
+  return module_create(instance, iree_allocator_system(), module);
 #endif
 }