Upload model

modeling.py

@@ -167,10 +167,12 @@ class TotalClassifierModel(PreTrainedModel):
             # inner list - list of above for each slice
             # outer list - list of above for each batch element (studies)
             batch_list = []
+            # probas.shape = (batch_size, num_slices, num_classes)
             for i in range(probas.shape[0]):
                 probas_i = probas[i]
+                # probas_i.shape = (num_slices, num_classes)
+                list_for_batch = []
                 for each_slice in range(probas_i.shape[0]):
-                    list_for_batch = []
                     for each_class in range(probas_i.shape[1]):
                         list_for_batch.append(
                             [
@@ -179,7 +181,7 @@ class TotalClassifierModel(PreTrainedModel):
                                 if probas_i[each_slice, each_class] >= threshold
                             ]
                         )
-                    batch_list.append(list_for_batch)
+                batch_list.append(list_for_batch)
             return batch_list
 
         return probas