Learnable embeddings are not cleaned up when using top-k model parameter (related to SageMaker HPO)

[thvasilo]

Currently, when a user sets a specific number of top-k models to be retained, the model files are removed but the learnable embeddings of every epoch are maintained.

The expected behavior is that only the embeddings for the topk best models should be retained.

Here's some example output from a SageMaker training job:

model_checkpoint $ tree
.
├── epoch-0
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-1
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-29
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-3
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-30
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-36
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-4
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-43
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-45
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-5
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt
├── epoch-51
│		 ├── author
│		 │		 ├── sparse_emb_00000.pt
│		 │		 └── sparse_emb_00001.pt
│		 ├── model.bin
│		 ├── optimizers.bin
│		 ├── paper
│		 │		 ├── sparse_emb_00000.pt
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     ├── sparse_emb_00000.pt
│		     └── sparse_emb_00001.pt
├── epoch-6
│		 ├── author
│		 │		 └── sparse_emb_00001.pt
│		 ├── paper
│		 │		 └── sparse_emb_00001.pt
│		 └── subject
│		     └── sparse_emb_00001.pt