This commit is contained in:
wataru 2023-04-14 15:19:09 +09:00
parent 8ddbb96500
commit 42da74d557

View File

@ -111,11 +111,15 @@ def export2onnx(input_model, output_model, output_model_simple, is_half):
except:
is_half = False
dev = torch.device("cpu")
net_g_onnx.eval().to(dev)
net_g_onnx.eval().to(dev).float()
net_g_onnx.load_state_dict(cpt["weight"], strict=False)
if is_half:
print("!!!!!!!!!!!!!!!!!! half")
net_g_onnx = net_g_onnx.half()
else:
print("!!!!!!!!!!!!!!!!!! full")
net_g_onnx = net_g_onnx.full()
if is_half:
feats = torch.HalfTensor(1, 2192, 256).to(dev)