fix csv logger

QData · Apr 24, 2020 · e5210e9 · e5210e9
1 parent fa059fb
commit e5210e9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/scripts/run_attack_args_helper.py b/scripts/run_attack_args_helper.py
@@ -278,9 +278,9 @@ def parse_logger_from_args(args):# Create logger
     # CSV
     if args.enable_csv:
         outfile_name = 'attack-{}.csv'.format(out_time)
-        plain = args.enable_csv == 'plain'
+        color_method = None if args.enable_csv == 'plain' else 'file'
         csv_path = os.path.join(args.out_dir, outfile_name)
-        attack_logger.add_output_csv(csv_path, plain)
+        attack_logger.add_output_csv(csv_path, color_method)
         print('Logging to CSV at path {}.'.format(csv_path))
 
     # Visdom

diff --git a/textattack/loggers/attack_logger.py b/textattack/loggers/attack_logger.py
@@ -23,8 +23,8 @@ def enable_visdom(self):
     def add_output_file(self, filename):
         self.loggers.append(FileLogger(filename=filename))
 
-    def add_output_csv(self, filename, plain):
-        self.loggers.append(CSVLogger(filename=filename, plain=plain))
+    def add_output_csv(self, filename, color_method):
+        self.loggers.append(CSVLogger(filename=filename, color_method=color_method))
 
     def log_result(self, result):
         self.results.append(result)