diff --git a/build/classes/java/main/controller/Controller.class b/build/classes/java/main/controller/Controller.class
index 54935a9eb3026ff1122cd50fbbcce457729b381e..d804fd9ea7c85c723b33fd680353024a4e86a2c5 100644
Binary files a/build/classes/java/main/controller/Controller.class and b/build/classes/java/main/controller/Controller.class differ
diff --git a/src/main/java/controller/Controller.java b/src/main/java/controller/Controller.java
index e3e7518cc2d1e9a99c7f4f99c30311759a926dab..25585b45bd5a53ae1857b402339c57baa28bc2e3 100644
--- a/src/main/java/controller/Controller.java
+++ b/src/main/java/controller/Controller.java
@@ -12,9 +12,7 @@ import javafx.scene.control.ToggleButton;
 import javafx.scene.control.ToggleGroup;
 import javafx.util.Duration;
 import javafx.util.Pair;
-import model.Board;
-import model.ModelElement;
-import model.FirefighterBoard;
+import model.*;
 import util.Position;
 import view.Grid;
 import view.ViewElement;
@@ -80,17 +78,19 @@ public class Controller {
     grid.repaint(viewElements);
     updateGenerationLabel(board.getBehavior().stepNumber());
   }
-
   private ViewElement getViewElement(List<ModelElement> squareState) {
-    if(squareState.contains(ModelElement.FIREFIGHTER)){
-      return ViewElement.FIREFIGHTER;
-    }
-    if (squareState.contains(ModelElement.FIRE)){
-      return ViewElement.FIRE;
+    for (ModelElement element : squareState) {
+      if (element instanceof FireFighter) {
+        return ViewElement.FIREFIGHTER;
+      }
+      if (element instanceof Fire) {
+        return ViewElement.FIRE;
+      }
     }
     return ViewElement.EMPTY;
   }
 
+
   private void initializeTimeline() {
     Duration duration = new Duration(Controller.PERIOD_IN_MILLISECONDS);
     EventHandler<ActionEvent> eventHandler =