Skip to content

GitLab

Unverified Commit f03ca0f9 authored by Vasilis Vryniotis's avatar Vasilis Vryniotis Committed by GitHub
Browse files

Switch torch.quantization to torch.ao.quantization (#5296)

parent 693a4632
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
torchvision/prototype/models/quantization/mobilenetv3.py
View file @ f03ca0f9
...@@ -43,14 +43,14 @@ def _mobilenet_v3_model( ...@@ -43,14 +43,14 @@ def _mobilenet_v3_model(
if quantize: if quantize:
model.fuse_model() model.fuse_model()
model.qconfig = torch.quantization.get_default_qat_qconfig(backend) model.qconfig = torch.ao.quantization.get_default_qat_qconfig(backend)
torch.quantization.prepare_qat(model, inplace=True) torch.ao.quantization.prepare_qat(model, inplace=True)
if weights is not None: if weights is not None:
model.load_state_dict(weights.get_state_dict(progress=progress)) model.load_state_dict(weights.get_state_dict(progress=progress))
if quantize: if quantize:
torch.quantization.convert(model, inplace=True) torch.ao.quantization.convert(model, inplace=True)
model.eval() model.eval()
return model return model
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment