ShadeNet / config.py
singam96's picture
Initial
12510fb
Raw
History Blame Contribute Delete
403 Bytes
import torch
CONFIG = {
"image_size": 512,
"width_mult": 2.0,
"device": "cuda" if torch.cuda.is_available() else "cpu",
"checkpoint_path": "./checkpoints/last.ckpt",
"onnx_dir": "./onnx",
"onnx_mode0": "model_mode0.onnx",
"onnx_mode0_quantized": "model_mode0_quantized.onnx",
"onnx_mode1": "model_mode1.onnx",
"onnx_mode1_quantized": "model_mode1_quantized.onnx",
}