From 0afcacea242a1982b86da54205d8011733cd1f5c Mon Sep 17 00:00:00 2001
From: Justin <605492+justinpinkney@users.noreply.github.com>
Date: Fri, 23 Sep 2022 09:21:53 +0100
Subject: [PATCH] Update slimify.py

---
 scripts/slimify.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/scripts/slimify.py b/scripts/slimify.py
index 7827135..3da87b0 100644
--- a/scripts/slimify.py
+++ b/scripts/slimify.py
@@ -1,9 +1,8 @@
 import torch
 import argparse
 
-# orig_path = "logs/2022-09-02T06-46-25_pokemon_pokemon/checkpoints/epoch=000142.ckpt"
-# out_name = "pokemon-ema-only.ckpt"
 if __name__ == "__main__":
+    # Make a version of the checkpoint with only ema weights (around 4GB)
     parser = argparse.ArgumentParser()
     parser.add_argument("--original_ckpt", help="full size checkpoint file")
     parser.add_argument("--output_path", help="filename for ema only checkpoint")
@@ -17,4 +16,4 @@ if __name__ == "__main__":
     new_d["state_dict"] = ema_state
 
     print(f"saving to {args.output_path}")
-    torch.save(new_d, args.output_path)
\ No newline at end of file
+    torch.save(new_d, args.output_path)