Make wandb checkpoint logging compatible with wandb model registry

composer/loggers/wandb_logger.py

@@ -239,6 +239,10 @@ def upload_file(self, state: State, remote_file_name: str, file_path: pathlib.Pa
                 # the trainer is evaluating or predicting. Assuming evaluation in this case.
                 metadata.update({f'eval_timestamp/{k}': v for (k, v) in state.eval_timestamp.state_dict().items()})
 
+            # Change the extension so the checkpoint is compatible with W&B's model registry
+            if extension == 'pt':
+                extension = 'model'
+
             wandb_artifact = wandb.Artifact(
                 name=new_remote_file_name,
                 type=extension,