Checkpoint

This commit is contained in:
Tibo De Peuter 2025-12-09 15:10:12 +01:00
parent 8311eabd4d
commit 926cde17d3
Signed by: tdpeuter
GPG key ID: 38297DE43F75FFE2
3 changed files with 24 additions and 13 deletions

24
main.py
View file

@ -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")