Enable stage-2 MMU.
diff --git a/src/api.c b/src/api.c
index 31ff376..0030473 100644
--- a/src/api.c
+++ b/src/api.c
@@ -49,7 +49,7 @@
 		return HF_VCPU_WAIT_FOR_INTERRUPT;
 	}
 
-	arch_set_vm_mm(&vm->page_table);
+	vm_set_current(vm);
 	*next = vcpu;
 
 	return HF_VCPU_YIELD;
@@ -64,7 +64,7 @@
 	struct vcpu *vcpu = &primary_vm.vcpus[cpu_index(cpu())];
 
 	/* Switch back to primary VM. */
-	arch_set_vm_mm(&primary_vm.page_table);
+	vm_set_current(&primary_vm);
 
 	/*
 	 * Inidicate to primary VM that this vcpu blocked waiting for an