diff --git a/prototypes/locomotive.lua b/prototypes/locomotive.lua
index e945663..b1af1e4 100644
--- a/prototypes/locomotive.lua
+++ b/prototypes/locomotive.lua
@@ -12,7 +12,7 @@ function locomotive:new(entity)
 end
 
 function locomotive:update()
-	if not self.entity.valid or not self.entity.is_headed_to_trains_front  then
+	if not self.entity.valid  then
 		return
 	end
 	local required_fuel=self.stack_size-self.entity.get_item_count()
diff --git a/prototypes/rail-power.lua b/prototypes/rail-power.lua
index cc1285f..ae23ea6 100644
--- a/prototypes/rail-power.lua
+++ b/prototypes/rail-power.lua
@@ -36,6 +36,9 @@ function railprototype:on_built()
 end
 
 function railprototype:connect()
+	if not self.entity or not self.entity.valid then
+		return
+	end
 	self.electricnode.disconnect_neighbour()
 	--connect to other rails
 	for _,rail in pairs(self:get_connected_powerrails()) do