diff --git a/plugins/UltimakerMachineActions/UMOUpgradeSelection.py b/plugins/UltimakerMachineActions/UMOUpgradeSelection.py
index 64c9ae1180..b92dc30c68 100644
--- a/plugins/UltimakerMachineActions/UMOUpgradeSelection.py
+++ b/plugins/UltimakerMachineActions/UMOUpgradeSelection.py
@@ -31,7 +31,7 @@ class UMOUpgradeSelection(MachineAction):
             if variant:
                 if variant.getId() == "empty_variant":
                     variant_index = global_container_stack.getContainerIndex(variant)
-                    self._createVariant(global_container_stack, variant_index)
+                    variant = self._createVariant(global_container_stack, variant_index)
                 variant.setProperty("machine_heated_bed", "value", heated_bed)
                 self.heatedBedChanged.emit()
 
@@ -41,4 +41,5 @@ class UMOUpgradeSelection(MachineAction):
         new_variant.addMetaDataEntry("type", "variant")
         new_variant.setDefinition(global_container_stack.getBottom())
         UM.Settings.ContainerRegistry.getInstance().addContainer(new_variant)
-        global_container_stack.replaceContainer(variant_index, new_variant)
\ No newline at end of file
+        global_container_stack.replaceContainer(variant_index, new_variant)
+        return new_variant
\ No newline at end of file