| # This script exports pre-trained model weights in the safetensors format. | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| from safetensors import torch as stt | |
| m = torchvision.models.efficientnet_b5(pretrained=True) | |
| stt.save_file(m.state_dict(), 'efficientnet-b5.safetensors') | |
| m = torchvision.models.efficientnet_b6(pretrained=True) | |
| stt.save_file(m.state_dict(), 'efficientnet-b6.safetensors') | |
| m = torchvision.models.efficientnet_b7(pretrained=True) | |
| stt.save_file(m.state_dict(), 'efficientnet-b7.safetensors') | |