File tree 1 file changed +6
-0
lines changed
1 file changed +6
-0
lines changed Original file line number Diff line number Diff line change @@ -20,6 +20,7 @@ pub(crate) fn is_pytorch(file_path: &Path) -> bool {
20
20
. to_ascii_lowercase ( ) ;
21
21
22
22
file_ext == "pt"
23
+ || file_ext == "pth"
23
24
|| file_name. ends_with ( "pytorch_model.bin" )
24
25
// cases like diffusion_pytorch_model.fp16.bin
25
26
|| ( file_name. contains ( "pytorch_model" ) && file_name. ends_with ( ".bin" ) )
@@ -54,6 +55,11 @@ mod tests {
54
55
assert ! ( is_pytorch( Path :: new( "path/to/model.pt" ) ) ) ;
55
56
assert ! ( is_pytorch( Path :: new( "MODEL.PT" ) ) ) ; // Case insensitive
56
57
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
+
57
63
// Standard pytorch_model.bin filename
58
64
assert ! ( is_pytorch( Path :: new( "pytorch_model.bin" ) ) ) ;
59
65
assert ! ( is_pytorch( Path :: new( "path/to/pytorch_model.bin" ) ) ) ;
You can’t perform that action at this time.
0 commit comments