Skip to content

Commit 7a0bd09

Browse files
committed
Update model prune loader to use pkgutil
1 parent 0f2803d commit 7a0bd09

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

timm/models/_prune.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import os
2+
import pkgutil
23
from copy import deepcopy
34

45
from torch import nn as nn
@@ -108,6 +109,5 @@ def adapt_model_from_string(parent_module, model_string):
108109

109110

110111
def adapt_model_from_file(parent_module, model_variant):
111-
adapt_file = os.path.join(os.path.dirname(__file__), '_pruned', model_variant + '.txt')
112-
with open(adapt_file, 'r') as f:
113-
return adapt_model_from_string(parent_module, f.read().strip())
112+
adapt_data = pkgutil.get_data(__name__, os.path.join('_pruned', model_variant + '.txt'))
113+
return adapt_model_from_string(parent_module, adapt_data.decode('utf-8').strip())

0 commit comments

Comments
 (0)