From 84ae9f95683412b8722ae56cf35ed2e999eef133 Mon Sep 17 00:00:00 2001
From: edward-io <53842584+edward-io@users.noreply.github.com>
Date: Wed, 9 Oct 2019 13:25:36 -0700
Subject: [PATCH] fix insights example (#102)

---
 captum/insights/example.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/captum/insights/example.py b/captum/insights/example.py
index 8dc8ad49b8..6606c8a3bb 100644
--- 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__":