diff --git a/src/main/java/li/cil/oc/server/driver/item/UpgradeGenerator.scala b/src/main/java/li/cil/oc/server/driver/item/UpgradeGenerator.scala
index 063a6ea90..91c3523d9 100644
--- a/src/main/java/li/cil/oc/server/driver/item/UpgradeGenerator.scala
+++ b/src/main/java/li/cil/oc/server/driver/item/UpgradeGenerator.scala
@@ -2,6 +2,7 @@ package li.cil.oc.server.driver.item
 
 import li.cil.oc.Items
 import li.cil.oc.api.driver.Slot
+import li.cil.oc.api.machine.Robot
 import li.cil.oc.server.component
 import net.minecraft.item.ItemStack
 import net.minecraft.tileentity.TileEntity
@@ -9,7 +10,9 @@ import net.minecraft.tileentity.TileEntity
 object UpgradeGenerator extends Item {
   override def worksWith(stack: ItemStack) = isOneOf(stack, Items.upgradeGenerator)
 
-  override def createEnvironment(stack: ItemStack, container: TileEntity) = new component.UpgradeGenerator(container)
+  override def createEnvironment(stack: ItemStack, container: TileEntity) = container match {
+    case robot: Robot => new component.UpgradeGenerator(robot)
+  }
 
   override def slot(stack: ItemStack) = Slot.Upgrade
 }