Fix collate_fn changing hash of torch trainer.

TeamEpochGithub · Jun 10, 2024 · f5a3b64 · f5a3b64
1 parent be8e21d
commit f5a3b64
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/epochalyst/pipeline/model/training/torch_trainer.py b/epochalyst/pipeline/model/training/torch_trainer.py
@@ -165,7 +165,7 @@ def log_to_terminal(self, message: str) -> None:
     epochs: Annotated[int, Gt(0)] = 10
     patience: Annotated[int, Gt(0)] = 5  # Early stopping
     batch_size: Annotated[int, Gt(0)] = 32
-    collate_fn: Callable[[tuple[Tensor, ...]], tuple[Tensor, ...]] = custom_collate
+    collate_fn: Callable[[tuple[Tensor, ...]], tuple[Tensor, ...]] = field(default=custom_collate, init=True, repr=False, compare=False)
 
     # Checkpointing
     checkpointing_enabled: bool = field(default=True, init=True, repr=False, compare=False)