From 7a0bd095cb184d702a208d4d05f06f4fa426a9fd Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Mon, 6 Feb 2023 17:45:16 -0800 Subject: [PATCH] Update model prune loader to use pkgutil --- timm/models/_prune.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/timm/models/_prune.py b/timm/models/_prune.py index 4e744dec..9bbe71ec 100644 --- a/timm/models/_prune.py +++ b/timm/models/_prune.py @@ -1,4 +1,5 @@ import os +import pkgutil from copy import deepcopy from torch import nn as nn @@ -108,6 +109,5 @@ def adapt_model_from_string(parent_module, model_string): def adapt_model_from_file(parent_module, model_variant): - adapt_file = os.path.join(os.path.dirname(__file__), '_pruned', model_variant + '.txt') - with open(adapt_file, 'r') as f: - return adapt_model_from_string(parent_module, f.read().strip()) + adapt_data = pkgutil.get_data(__name__, os.path.join('_pruned', model_variant + '.txt')) + return adapt_model_from_string(parent_module, adapt_data.decode('utf-8').strip())