diff --git a/src/deep_neurographs/inference.py b/src/deep_neurographs/inference.py
index de946664..630d7a2f 100644
--- a/src/deep_neurographs/inference.py
+++ b/src/deep_neurographs/inference.py
@@ -4,14 +4,14 @@
 @author: Anna Grim
 @email: anna.grim@alleninstitute.org
 
-Routines for running inference with machine models that classifies edge
-proposals.
+Routines for running inference with a machine model that classifies edge proposals.
 
 """
 
 from datetime import datetime
 from time import time
 from torch.nn.functional import sigmoid
+from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 import networkx as nx
diff --git a/src/deep_neurographs/utils/util.py b/src/deep_neurographs/utils/util.py
index e8c9a68f..023b6a6c 100644
--- a/src/deep_neurographs/utils/util.py
+++ b/src/deep_neurographs/utils/util.py
@@ -372,7 +372,7 @@ def write_txt(path, contents):
     f.close()
 
 
-def write_to_s3(local_path, bucket_name, s3_key):
+def write_to_s3(local_path, bucket_name, prefix):
     """
     Writes a single file on local machine to an s3 bucket.
 
@@ -382,7 +382,7 @@ def write_to_s3(local_path, bucket_name, s3_key):
         Path to file to be written to s3.
     bucket_name : str
         Name of s3 bucket.
-    s3_key : str
+    prefix : str
         Path within s3 bucket.
 
     Returns
@@ -391,7 +391,7 @@ def write_to_s3(local_path, bucket_name, s3_key):
 
     """
     s3 = boto3.client('s3')
-    s3.upload_file(local_path, bucket_name, s3_key)
+    s3.upload_file(local_path, bucket_name, prefix)
 
 
 # --- math utils ---