fixing save-all/save-recycles option(s)

Author: sokrypton

colabfold/batch.py

@@ -428,23 +428,20 @@ def callback(prediction_result, recycles):
                     print_line += f" {y}={prediction_result[x]:.3g}"
                 logger.info(f"{tag} recycle={recycles}{print_line}")
 
-                if save_recycles or save_all:
-                    prediction_result = _jnp_to_np(prediction_result)
-                    prediction_result["representations"] = prediction_result.pop("prev")
-                
                 if save_recycles:
-                    final_atom_mask = prediction_result["structure_module"]["final_atom_mask"]
-                    b_factors = prediction_result["plddt"][:, None] * final_atom_mask
+                    result = _jnp_to_np(prediction_result)
+                    final_atom_mask = result["structure_module"]["final_atom_mask"]
+                    b_factors = result["plddt"][:, None] * final_atom_mask
                     unrelaxed_protein = protein.from_prediction(features=input_features,
-                        result=prediction_result, b_factors=b_factors,
+                        result=result, b_factors=b_factors,
                         remove_leading_feature_dimension=("ptm" in model_type))
 
                     unrelaxed_pdb_lines = protein.to_pdb(class_to_np(unrelaxed_protein))
                     files.get("unrelaxed",f"r{recycles}.pdb").write_text(unrelaxed_pdb_lines)
 
-                if save_all:
-                    with files.get("all",f"r{recycles}.pickle").open("wb") as handle:
-                        pickle.dump(prediction_result, handle)
+                    if save_all:
+                        with files.get("all",f"r{recycles}.pickle").open("wb") as handle:
+                            pickle.dump(result, handle)
 
             prediction_result, recycles = \
             model_runner.predict(input_features, random_seed=seed, prediction_callback=callback)
@@ -489,7 +486,11 @@ def callback(prediction_result, recycles):
             files.get("unrelaxed","pdb").write_text(protein_lines)
             unrelaxed_pdb_lines.append(protein_lines)
 
+
             # save raw outputs
+            if save_all:
+                with files.get("all","pickle").open("wb") as handle:
+                    pickle.dump(prediction_result, handle)
             if save_single_representations or save_pair_representations:
                 rep = prediction_result["representations"]
                 if save_single_representations: