Fix bug

src/femr/models/tasks.py

@@ -296,7 +296,10 @@ def fit_pretraining_task_info(
         for task, task_stats in zip(tasks, stats):
             frac_events = task_stats[1] / (task_stats[0] + task_stats[1])
             rate = frac_events / task_stats[2].mean()
-            task_data.append((task, rate))
+            if rate != 0:
+                task_data.append((task, rate))
+            else:
+                print("Ran into task of rate 0?", task, frac_events, task_stats[0], task_stats[1], task_stats[2].mean())        
 
         return MOTORTask(task_data, time_bins, final_layer_size)
 

src/femr/models/transformer.py

@@ -364,7 +364,7 @@ def compute_features(
     task = femr.models.tasks.LabeledSubjectTask(labels)
 
     model = femr.models.transformer.FEMRModel.from_pretrained(model_path, task_config=task.get_task_config())
-    tokenizer = femr.models.tokenizer.FEMRTokenizer.from_pretrained(model_path, ontology=ontology)
+    tokenizer = femr.models.tokenizer.HierarchicalTokenizer.from_pretrained(model_path, ontology=ontology)
     processor = femr.models.processor.FEMRBatchProcessor(tokenizer, task=task)
 
     filtered_data = db.filter(list(task.label_map.keys()))
@@ -396,6 +396,8 @@ def compute_features(
                 all_feature_times.append(result["timestamps"].to(cpu_device, non_blocking=True))
                 all_representations.append(result["representations"].to(cpu_device, non_blocking=True))
 
+    torch.cuda.synchronize()
+
     all_subject_ids_np = torch.concatenate(all_subject_ids).numpy()
     all_feature_times_np = torch.concatenate(all_feature_times).numpy()
     all_representations_np = torch.concatenate(all_representations).numpy()