diff --git a/lib/puppet_litmus/rake_tasks.rb b/lib/puppet_litmus/rake_tasks.rb
index dcb2fe6..3f08a59 100644
--- a/lib/puppet_litmus/rake_tasks.rb
+++ b/lib/puppet_litmus/rake_tasks.rb
@@ -128,6 +128,7 @@
     Rake::Task['spec_prep'].invoke
 
     results = install_agent(args[:collection], targets, inventory_hash)
+    target_index = 0
     results.each do |result|
       command_to_run = "bolt task run puppet_agent::install --targets #{result['target']} --inventoryfile spec/fixtures/litmus_inventory.yaml --modulepath #{DEFAULT_CONFIG_DATA['modulepath']}"
       raise "Failed on #{result['target']}\n#{result}\ntry running '#{command_to_run}'" if result['status'] != 'success'
@@ -157,7 +158,8 @@
       end
 
       # add puppet-agent feature to successful nodes
-      inventory_hash = add_feature_to_node(inventory_hash, 'puppet-agent', result['target'])
+      inventory_hash = add_feature_to_node(inventory_hash, 'puppet-agent', targets[target_index])
+      target_index += 1
     end
 
     # update the inventory with the puppet-agent feature set per node