diff --git a/dan/predict/prediction.py b/dan/predict/prediction.py
index d3e7a5fe646b64c95cf81d8ac800065956a3da0b..8d021712c2d36d049566f81ebbee5f78021a4e16 100644
--- a/dan/predict/prediction.py
+++ b/dan/predict/prediction.py
@@ -105,7 +105,6 @@ class DAN:
         start_token=None,
         threshold_method="otsu",
         threshold_value=0,
-        temperature=1.0,
     ):
         """
         Run prediction on an input image.
@@ -179,7 +178,7 @@ class DAN:
                     num_pred=1,
                 )
 
-                pred = pred / temperature
+                pred = pred / self.temperature
                 whole_output.append(output)
                 attention_maps.append(weights)
                 confidence_scores.append(