diff --git a/open_flamingo/train/data.py b/open_flamingo/train/data.py
index 7d46fd14..ab744c60 100644
--- a/open_flamingo/train/data.py
+++ b/open_flamingo/train/data.py
@@ -14,7 +14,7 @@
 import webdataset as wds
 from PIL import Image
 import base64
-from scipy import linear_sum_assignment
+from scipy.optimize import linear_sum_assignment
 
 from data_utils import *