Skip to content

Commit 5ad824a

Browse files
committed
new: added .pth support
1 parent a8fa2fe commit 5ad824a

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

src/core/pytorch/mod.rs

+6
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ pub(crate) fn is_pytorch(file_path: &Path) -> bool {
2020
.to_ascii_lowercase();
2121

2222
file_ext == "pt"
23+
|| file_ext == "pth"
2324
|| file_name.ends_with("pytorch_model.bin")
2425
// cases like diffusion_pytorch_model.fp16.bin
2526
|| (file_name.contains("pytorch_model") && file_name.ends_with(".bin"))
@@ -54,6 +55,11 @@ mod tests {
5455
assert!(is_pytorch(Path::new("path/to/model.pt")));
5556
assert!(is_pytorch(Path::new("MODEL.PT"))); // Case insensitive
5657

58+
// Standard .pth extension
59+
assert!(is_pytorch(Path::new("model.pth")));
60+
assert!(is_pytorch(Path::new("path/to/model.pth")));
61+
assert!(is_pytorch(Path::new("MODEL.PTH"))); // Case insensitive
62+
5763
// Standard pytorch_model.bin filename
5864
assert!(is_pytorch(Path::new("pytorch_model.bin")));
5965
assert!(is_pytorch(Path::new("path/to/pytorch_model.bin")));

0 commit comments

Comments
 (0)