diff --git a/action_plugins/qubesformation.py b/action_plugins/qubesformation.py
index a386ded..da54f08 100644
--- a/action_plugins/qubesformation.py
+++ b/action_plugins/qubesformation.py
@@ -17,12 +17,13 @@ def generate_datastructure(vms, task_vars):
     dc = collections.OrderedDict
     d = dc()
     for n, data in vms.items():
-        # This block will skip any VMs that are not in the 'vms' variable
+        # This block will skip any VMs that are not in the groups defined in the 'formation_vm_groups' variable
         # This allows you to deploy in multiple stages which is useful in cases
         # where you want to create a template after another template is already provisioned.
-        if hasattr(task_vars, 'vms'):
-            if n not in task_vars['vms']:
-                continue
+        if hasattr(task_vars, 'formation_vm_groups'):
+            for group in task_vars['formation_vm_groups']:
+                if n not in task_vars.groups[group]:
+                    continue
         
         qubes = data['qubes']
         d[task_vars['hostvars'][n]['inventory_hostname_short']] = dc(qvm=['vm'])