diff --git a/kg_env.py b/kg_env.py
index 9f85f61..5e57a78 100644
--- a/kg_env.py
+++ b/kg_env.py
@@ -222,7 +222,6 @@ def batch_step(self, batch_act_idx):
             act_idx = batch_act_idx[i]
             _, curr_node_type, curr_node_id = self._batch_path[i][-1]
             relation, next_node_id = self._batch_curr_actions[i][act_idx]
-            # TODO: check act_idx is valid?
             if relation == SELF_LOOP:
                 next_node_type = curr_node_type
             else:
diff --git a/knowledge_graph.py b/knowledge_graph.py
index b61c4a6..7ee377b 100644
--- a/knowledge_graph.py
+++ b/knowledge_graph.py
@@ -204,23 +204,6 @@ def heuristic_search(self, uid, pid, pattern_id, trim_edges=False):
                 intersect_nodes = wids_u.intersection(wids_u_p)
                 tmp_paths = [(uid, x, uid_p, pid) for x in intersect_nodes]
                 paths.extend(tmp_paths)
-        # elif len(pattern) == 5:  # DOES NOT WORK SO FAR!
-        #    nodes_from_user = set(self.G[USER][uid][pattern[1][0]])  # USER->MENTION->WORD
-        #    nodes_from_product = set(self.G[PRODUCT][pid][pattern[-1][0]])
-        #    if pattern[-2][1] == USER:
-        #        nodes_from_product.difference([uid])
-        #    count = 0
-        #    for wid in nodes_from_user:
-        #        pids_from_wid = set(self.G[WORD][wid][pattern[2][0]])  # USER->MENTION->WORD->DESCRIBE->PRODUCT
-        #        pids_from_wid = pids_from_wid.difference([pid])  # exclude target product
-        #        for nid in nodes_from_product:
-        #            if pattern[-2][1] == WORD:
-        #                if nid == wid:
-        #                    continue
-        #            other_pids = set(self.G[pattern[-2][1]][nid][pattern[-2][0]])
-        #            intersect_nodes = pids_from_wid.intersection(other_pids)
-        #            count += len(intersect_nodes)
-        #    return count
 
         return paths
 
@@ -236,3 +219,4 @@ def check_test_path(dataset_str, kg):
                 count += len(tmp_path)
             if count == 0:
                 print(uid, pid)
+
diff --git a/preprocess.py b/preprocess.py
index 885d228..3eb145d 100644
--- a/preprocess.py
+++ b/preprocess.py
@@ -27,7 +27,7 @@ def generate_labels(dataset, mode='train'):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset', type=str, default=CELL, help='One of {BEAUTY, CELL, CD, CLOTH}.')
+    parser.add_argument('--dataset', type=str, default=BEAUTY, help='One of {BEAUTY, CELL, CD, CLOTH}.')
     args = parser.parse_args()
 
     # Create AmazonDataset instance for dataset.
@@ -55,7 +55,6 @@ def main():
     # =========== END =========== #
 
 
-
 if __name__ == '__main__':
     main()
 
diff --git a/test_agent.py b/test_agent.py
index 6d1ef9e..9ace7f9 100644
--- a/test_agent.py
+++ b/test_agent.py
@@ -13,7 +13,6 @@
 import torch.optim as optim
 from torch.autograd import Variable
 from torch.distributions import Categorical
-from tensorboardX import SummaryWriter
 import threading
 from functools import reduce
 
@@ -220,7 +219,7 @@ def test(args):
 if __name__ == '__main__':
     boolean = lambda x: (str(x).lower() == 'true')
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset', type=str, default='beauty', help='One of {cloth, beauty, cell, cd}')
+    parser.add_argument('--dataset', type=str, default=BEAUTY, help='One of {cloth, beauty, cell, cd}')
     parser.add_argument('--name', type=str, default='train_agent', help='directory name.')
     parser.add_argument('--seed', type=int, default=123, help='random seed.')
     parser.add_argument('--gpu', type=str, default='0', help='gpu device.')
diff --git a/train_agent.py b/train_agent.py
index 8ab039e..5786540 100644
--- a/train_agent.py
+++ b/train_agent.py
@@ -10,7 +10,6 @@
 import torch.optim as optim
 from torch.autograd import Variable
 from torch.distributions import Categorical
-from tensorboardX import SummaryWriter
 
 from knowledge_graph import KnowledgeGraph
 from kg_env import BatchKGEnvironment
@@ -130,8 +129,6 @@ def get_batch(self):
 
 
 def train(args):
-    train_writer = SummaryWriter(args.log_dir)
-
     env = BatchKGEnvironment(args.dataset, args.max_acts, max_path_len=args.max_path_len, state_history=args.state_history)
     uids = list(env.kg(USER).keys())
     dataloader = ACDataLoader(uids, args.batch_size)
@@ -185,11 +182,6 @@ def train(args):
                         ' | vloss={:.5f}'.format(avg_vloss) +
                         ' | entropy={:.5f}'.format(avg_entropy) +
                         ' | reward={:.5f}'.format(avg_reward))
-                train_writer.add_scalar('train/loss', avg_loss, step)
-                train_writer.add_scalar('train/ploss', avg_ploss, step)
-                train_writer.add_scalar('train/vloss', avg_vloss, step)
-                train_writer.add_scalar('train/entropy', avg_entropy, step)
-                train_writer.add_scalar('train/reward', avg_reward, step)
         ### END of epoch ###
 
         policy_file = '{}/policy_model_epoch_{}.ckpt'.format(args.log_dir, epoch)
@@ -199,7 +191,7 @@ def train(args):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dataset', type=str, default='beauty', help='One of {clothing, cell, beauty, cd}')
+    parser.add_argument('--dataset', type=str, default=BEAUTY, help='One of {clothing, cell, beauty, cd}')
     parser.add_argument('--name', type=str, default='train_agent', help='directory name.')
     parser.add_argument('--seed', type=int, default=123, help='random seed.')
     parser.add_argument('--gpu', type=str, default='0', help='gpu device.')