Checkpoint
This commit is contained in:
parent
8311eabd4d
commit
926cde17d3
3 changed files with 24 additions and 13 deletions
24
main.py
24
main.py
|
|
@ -17,23 +17,27 @@ def main():
|
|||
match args.mode:
|
||||
case 'train':
|
||||
train(
|
||||
device = device,
|
||||
dataset = args.dataset,
|
||||
data_root = args.data_root,
|
||||
n_trials = 3 if args.debug else None,
|
||||
size = 2**10 if args.debug else None,
|
||||
method = args.method,
|
||||
device=device,
|
||||
dataset=args.dataset,
|
||||
data_root=args.data_root,
|
||||
n_trials=3 if args.debug else None,
|
||||
size=2 ** 10 if args.debug else None,
|
||||
method=args.method,
|
||||
model_name=args.model,
|
||||
model_path = args.model_load_path,
|
||||
model_out = args.model_save_path
|
||||
model_path=args.model_load_path,
|
||||
model_out=args.model_save_path
|
||||
)
|
||||
|
||||
case 'compress':
|
||||
compress(args.input_file)
|
||||
compress(device=device,
|
||||
model_path=args.model_load_path,
|
||||
input_file=args.input_file,
|
||||
output_file=args.output_file
|
||||
)
|
||||
|
||||
case _:
|
||||
raise NotImplementedError(f"Mode {args.mode} is not implemented yet")
|
||||
|
||||
|
||||
print("Done")
|
||||
|
||||
|
||||
|
|
|
|||
Reference in a new issue