diff --git a/dan/predict/attention.py b/dan/predict/attention.py
index 09229ec041cc2a6e314ad53de80ef6b6784fc762..91ae590e80d05e3ac12e268098c77e98f912e12e 100644
--- a/dan/predict/attention.py
+++ b/dan/predict/attention.py
@@ -3,7 +3,6 @@ import re
 
 import cv2
 import numpy as np
-import math
 from PIL import Image
 
 from dan import logger
@@ -74,8 +73,8 @@ def plot_attention(
     :param scale: Scaling factor for the output gif image
     :param outname: Name of the gif image
     """
-    
-    height, width, _ = image.shape 
+
+    height, width, _ = image.shape
     attention_map = []
 
     # Convert to PIL Image and create mask