# Load pytorch trained model via optimum
**URL:**
**Author:**  [@Talha](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FTalha)%5C
**Post date:** [June 14, 2022, 6:13pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fload-pytorch-trained-model-via-optimum%2F19157%2F1 "2022-06-14T18:13:54Z")
I train a bert model using pytorch lightning now i want to load it to optimum for inference. How can i do that.
I tried to save it as
```auto
torch.save(model.bertmodel.state_dict(), 'bert.pth')
```
then try to load in optimum as
```auto
# The type of quantization to apply
qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
quantizer = ORTQuantizer.from_pretrained('bert.pth', feature="sequence-classification")
# Quantize the model!
quantizer.export(
onnx_model_path="model.onnx",
onnx_quantized_model_output_path="model-quantized.onnx",
quantization_config=qconfig,
)
```
the error it throw is
```auto
OSError: bert.pth is not a local folder and is not a valid model identifier listed on 'https://huggingface.co/models'
```
Is there a way to handle this without uploading model to huggingface?
---
**Author:**  [@philschmid](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fphilschmid)%5C
**Post date:** [June 15, 2022, 8:51am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fload-pytorch-trained-model-via-optimum%2F19157%2F2 "2022-06-15T08:51:22Z")
@Talha can you load your model weights with the `transformers` `AutoModelForXXX` classes?
If you can pass the “directory” of where the `pytorch_model.bin` is stored to the `from_pretrained` method but this directory also needs to current include files for the tokenizer as well.
We are currently working on removing this `transformers` dependency to make it easier to use with raw onnx checkpoints, but there for you also need a `model.onnx` file looking at your snippet it seems that you currently only have a pytorch checkpoint
---
**Author:**  [@Talha](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FTalha)%5C
**Post date:** [June 15, 2022, 5:03pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fload-pytorch-trained-model-via-optimum%2F19157%2F3 "2022-06-15T17:03:28Z")
currently i am saving it via
```auto
tokenizer.save_pretrained("model")
model.bertmodel.save_pretrained("model")
```
and its working, but its not allowing me to do modification in last layer of network where i just want to change number of output neurons
EDIT:
I was wrong, it is not saving my pre-trained model but the original model from hugging face
So i saved it this way
```auto
import os
tokenizer=AutoTokenizer.from_pretrained(model_name,local_files_only=False)
output_dir = "./model/"
# Step 1: Save a model, configuration and vocabulary that you have fine-tuned
# If we have a distributed model, save only the encapsulated model
# (it was wrapped in PyTorch DistributedDataParallel or DataParallel)
model_to_save = model.module if hasattr(model, 'module') else model
# If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(output_dir, 'pytorch_model.bin')
output_config_file = os.path.join(output_dir, 'config.json')
torch.save(model.bertmodel.state_dict(), output_model_file)
model.bertmodel.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(output_dir)
# load again
#Example for a Bert model
model = AutoModelForSequenceClassification.from_pretrained(output_dir,num_labels=1)
tokenizer = AutoTokenizer.from_pretrained(output_dir)
```
but when i try reading. i got this error
```auto
# The type of quantization to apply
qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
quantizer = ORTQuantizer.from_pretrained("model", feature="sequence-classification")
# Quantize the model!
quantizer.export(
onnx_model_path="model.onnx",
onnx_quantized_model_output_path="model-quantized.onnx",
quantization_config=qconfig,
)
```
Error
RuntimeError: Error(s) in loading state\_dict for DistilBertForSequenceClassification:
size mismatch for classifier.weight: copying a param with shape torch.Size([1, 768]) from checkpoint, the shape in current model is torch.Size([2, 768]).
size mismatch for classifier.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([2]).
I am looking a way to define number of labels in `ORTQuantizer.from_pretrained`
---
**Author:**  [@philschmid](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fphilschmid)%5C
**Post date:** [June 20, 2022, 12:06pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fload-pytorch-trained-model-via-optimum%2F19157%2F4 "2022-06-20T12:06:39Z")
can you run inference on your `model` it seems that you have a miss-match between training inference when saving the model? Are you sure you are saving the model with the classification head? it looks more like you are saving the BERT model without it.
You don’t need to define `num_labels` after you trained your models. Since those are needed to initialize a classification head for fine-tuning but this should already exist in your model.
My suggestion would be to take a look at how to fine-tune it with `transformers` and `Trainer` rather then some custom logic, which might not be 100% compatible. Here is an example: [notebooks/text\_classification.ipynb at main · huggingface/notebooks · GitHub](https://github.com/huggingface/notebooks/blob/main/examples/text_classification.ipynb)
---
**Author:**  [@shubhamgoel27](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fshubhamgoel27)%5C
**Post date:** [August 8, 2022, 8:23pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fload-pytorch-trained-model-via-optimum%2F19157%2F5 "2022-08-08T20:23:05Z")
Hi @philschmid
Are there any updates to this? I’m trying to load a Pytorch model stored locally with Optimum, but facing a lot of issues.
---
**Author:**  [@philschmid](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fphilschmid)%5C
**Post date:** [August 10, 2022, 7:17am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fload-pytorch-trained-model-via-optimum%2F19157%2F6 "2022-08-10T07:17:09Z")
Yes, this [PR](https://github.com/huggingface/optimum/pull/270) has been merged which allows to use any `model.onnx` with the `ORTQuantizer`. Documentation: [Quantization](https://huggingface.co/docs/optimum/main/en/onnxruntime/quantization)
But it is not yet release so you have to install from main