diff --git a/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java b/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java
index e081a3531..aebc9f494 100644
--- a/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java
+++ b/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java
@@ -348,12 +348,10 @@ public class DPSGeneratorEntity extends LivingEntity implements GeoEntity {
 
     public void beastCharge() {
         this.entityData.set(LEVEL, 7);
-        this.getEnergy().ifPresent(cap -> {
-            if (cap instanceof SyncedEntityEnergyStorage storage) {
-                storage.setCapacity(this.getMaxEnergy());
-                storage.setMaxExtract(this.getMaxTransfer());
-                storage.setEnergy(this.getMaxEnergy());
-            }
-        });
+        if (this.getCapability(Capabilities.EnergyStorage.ENTITY, null) instanceof SyncedEntityEnergyStorage storage) {
+            storage.setCapacity(this.getMaxEnergy());
+            storage.setMaxExtract(this.getMaxTransfer());
+            storage.setEnergy(this.getMaxEnergy());
+        }
     }
 }