minor editing

inference.py

@@ -18,7 +18,7 @@ def infer(model, input_file):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="Inference")
-    parser.add_argument('--svs_dir', type=str, default='./samples_CRC', help="")
+    parser.add_argument('--svs_dir', type=str, default='./samples', help="")
 
     args = parser.parse_args()
 

models/aggregator.py

@@ -6,23 +6,6 @@ from functools import partial
 
 from models.transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer
 
-
-# @dataclass
-# class VisionCfg:
-#     layers: Union[Tuple[int, int, int, int], int] = 6
-#     width: int = 512
-#     head_width: int = 64
-#     mlp_ratio: float = 4.0
-
-#     ls_init_value: Optional[float] = None  # layer scale initial value
-#     patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
-#     no_ln_pre: bool = False  # disable pre transformer LayerNorm
-#     pool_type: str = 'none'
-#     final_ln_after_pool: bool = True  # apply final LayerNorm after pooling
-#     output_tokens: bool = False
-#     act_kwargs: Optional[dict] = None
-#     norm_kwargs: Optional[dict] = None
-
 @dataclass
 class CLIPVisionCfg:
     layers: Union[Tuple[int, int, int, int], int] = 6

models/transformer.py

@@ -215,7 +215,6 @@ class ResidualAttentionBlock(nn.Module):
         self.ln_1 = norm_layer(d_model)
         
         if use_flex:
-            print("Flex_Attention!")
             self.attn = Flex_Attention(dim = d_model, num_heads=n_head, proj_drop=dropout, use_rel_bias=use_rel_bias)
         else:
             self.attn = Attention(dim = d_model, num_heads=n_head, batch_first=batch_first, proj_drop=dropout, attn_drop=dropout)