key error when using finetune model to infer

[nuanxinqing]

i finetuned model use custom data, the finetune code:

`#!/bin/bash now=$(date +"%Y%m%d_%H%M%S")

epoch=120 bs=2 gpus=2 lr=0.000005 encoder=vits dataset=custom # vkitti img_size=518 min_depth=0.001 max_depth=80 #20 # 80 for virtual kitti pretrained_from=../checkpoints/depth_anythingv2${encoder}.pth save_path=exp/custom #exp/hypersim # exp/vkitti master_port=20597 mkdir -p $save_path

CUDA_VISIBLE_DEVICES=1,2 python3 -m torch.distributed.launch \ --nproc_per_node=$gpus \ --nnodes 1 \ --node_rank=0 \ --master_addr=localhost \ --master_port=$master_port \ train.py --epoch $epoch --encoder $encoder --bs $bs --lr $lr --save-path $save_path --dataset $dataset \ --img-size $img_size --min-depth $min_depth --max-depth $max_depth --pretrained-from $pretrained_from \ --port $master_port 2>&1 | tee -a $save_path/$now.log `

while doing inference i met the model loading error:

and my infer code according to run.py `**model_configs = { 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}, 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]} }

encoder = 'vitl' # or 'vits', 'vitb' dataset = 'vkitti' # 'hypersim' for indoor model, 'vkitti' for outdoor model max_depth = 80 # 20 for indoor model, 80 for outdoor model finetune_custom_model_path =osp.join(f'metric_depth/exp/custom-'+f'{encoder}', 'latest.pth') model = DepthAnythingV2(**{model_configs[encoder], 'max_depth': max_depth}) model.load_state_dict(torch.load(f'{finetune_custom_model_path}', map_location='cpu'))`

how to correct my code? thank u

[LiheYoung]

Try this in the last line:

model.load_state_dict(torch.load(f'{finetune_custom_model_path}', map_location='cpu')['model'])

[nuanxinqing]

@LiheYoung thank you , i have solved