fix insights example (#102)

pytorch · Oct 9, 2019 · 84ae9f9 · 84ae9f9
1 parent 85d4d69
commit 84ae9f9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/captum/insights/example.py b/captum/insights/example.py
@@ -1,6 +1,6 @@
 import os
 
-from captum.insights.api import AttributionVisualizer, Data
+from captum.insights import AttributionVisualizer, Batch
 from captum.insights.features import ImageFeature
 
 import torch
@@ -71,7 +71,7 @@ def formatted_data_iter():
     )
     while True:
         images, labels = next(dataloader)
-        yield Data(inputs=images, labels=labels)
+        yield Batch(inputs=images, labels=labels)
 
 
 if __name__ == "__main__":