Add files using upload-large-folder tool

b2cb4a0 verified 3 months ago

27.6 kB

#
# For licensing see accompanying LICENSE file.
# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
#

import os
import shutil
import copy
import numpy as np
from pathlib import Path
from einops import repeat

import torch
import torch.nn.functional as F
from torch.optim.optimizer import Optimizer
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel
from torch.optim.swa_utils import AveragedModel
from torch.nn.utils import clip_grad_norm_

import lightning
import lightning.pytorch as pl
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
from fairscale.nn.wrap import enable_wrap, wrap

from onescience.utils.simplefold.esm_utils import _af2_to_esm, esm_registry
from onescience.datapipes.boltz_data_pipeline.types import Record, Structure
from onescience.utils.simplefold.boltz_utils import (
    weighted_rigid_align, 
    center_random_augmentation,
    process_structure, 
    save_structure
)


def logit_normal_sample(n=1, m=0.0, s=1.0):
    # Logit-Normal Sampling from https://arxiv.org/pdf/2403.03206.pdf
    u = torch.randn(n) * s + m
    t = 1 / (1 + torch.exp(-u))
    return t


def lddt_dist(dmat_predicted, dmat_true, mask, cutoff=15.0, per_atom=False):
    # NOTE: the mask is a pairwise mask which should have the identity elements already masked out
    # Compute mask over distances
    dists_to_score = (dmat_true < cutoff).float() * mask
    dist_l1 = torch.abs(dmat_true - dmat_predicted)

    score = 0.25 * (
        (dist_l1 < 0.5).float()
        + (dist_l1 < 1.0).float()
        + (dist_l1 < 2.0).float()
        + (dist_l1 < 4.0).float()
    )

    # Normalize over the appropriate axes.
    if per_atom:
        mask_no_match = torch.sum(dists_to_score, dim=-1) != 0
        norm = 1.0 / (1e-10 + torch.sum(dists_to_score, dim=-1))
        score = norm * (1e-10 + torch.sum(dists_to_score * score, dim=-1))
        return score, mask_no_match.float()
    else:
        norm = 1.0 / (1e-10 + torch.sum(dists_to_score, dim=(-2, -1)))
        score = norm * (1e-10 + torch.sum(dists_to_score * score, dim=(-2, -1)))
        total = torch.sum(dists_to_score, dim=(-1, -2))
        return score, total


class SimpleFold(pl.LightningModule):
    def __init__(
        self,
        architecture,
        processor,
        loss,
        path,
        sampler,
        optimizer=None,
        scheduler=None,
        plddt_module=None,
        ema_decay=0.999,
        esm_model="esm2_3B",
        aa_bolt_link=None,
        use_rigid_align=True,
        smooth_lddt_loss_weight=1.0,
        lddt_cutoff=15.0,
        clip_grad_norm_val=None,
        lddt_weight_schedule=False,
        plddt_training=False,
        sample_dir='artifacts/',
    ):
        super().__init__()
        self.save_hyperparameters(logger=False)

        self.model = architecture
        self.model_ema = AveragedModel(
            self.model,
            multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(
                self.hparams.ema_decay
            ),
            use_buffers=True,
        )
        self.model_ema.eval()

        self.loss = loss
        self.path = path
        self.sampler = sampler

        self.use_rigid_align = use_rigid_align
        self.lddt_cutoff = lddt_cutoff
        self.smooth_lddt_loss_weight = smooth_lddt_loss_weight
        self.use_smooth_lddt_loss = smooth_lddt_loss_weight > 0.0
        self.lddt_weight_schedule = lddt_weight_schedule
        self.plddt_training = plddt_training
        self.sample_dir = sample_dir

        self.aa_bolt_link = aa_bolt_link
        self.nval_steps = 0

        try:
            self.t_eps = self.sampler.t_eps
        except AttributeError:
            self.t_eps = 0.0

        self.use_esm = esm_model is not None
        if self.use_esm:
            self.esm_model, self.esm_dict = esm_registry[esm_model]()
            self.esm_model.eval()
            self.af2_to_esm = _af2_to_esm(self.esm_dict)
            print(f"Using ESM model: {esm_model}")
        else:
            self.esm_model = None
            self.esm_dict = None
            self.af2_to_esm = None

        self.plddt_module = plddt_module
        if self.plddt_training:
            assert self.plddt_module is not None, "PLDDT module must be provided for PLDDT training"
            self.model.eval()

    def register(self, name, tensor):
        self.register_buffer(name, tensor.type(torch.float32))

    def loss_masking(self, loss, atom_mask):
        loss_mask = repeat(atom_mask, "b s -> b s d", d=loss.shape[-1])
        loss *= loss_mask

        denom = torch.sum(atom_mask, -1, keepdim=True)
        denom = denom.unsqueeze(-1)
        loss = torch.sum(loss, dim=1, keepdim=True) / denom
        return loss

    def smooth_lddt_loss(
        self, 
        pred_coords,
        true_coords,
        # is_nucleotide,
        coords_mask,
        t,
    ):
        """Compute weighted alignment.
        Parameters
        ----------
        pred_coords: torch.Tensor
            The predicted atom coordinates
        true_coords: torch.Tensor
            The ground truth atom coordinates
        coords_mask: torch.Tensor
            The atoms mask
        """
        B, N, _ = true_coords.shape
        true_dists = torch.cdist(true_coords, true_coords)

        mask = (true_dists < self.lddt_cutoff).float()
        mask = mask * (1 - torch.eye(pred_coords.shape[1], device=pred_coords.device))
        mask = mask * (coords_mask.unsqueeze(-1) * coords_mask.unsqueeze(-2))

        # Compute distances between all pairs of atoms
        pred_dists = torch.cdist(pred_coords, pred_coords)
        dist_diff = torch.abs(true_dists - pred_dists)

        # Compute epsilon values
        eps = (
            (
                (
                    F.sigmoid(0.5 - dist_diff)
                    + F.sigmoid(1.0 - dist_diff)
                    + F.sigmoid(2.0 - dist_diff)
                    + F.sigmoid(4.0 - dist_diff)
                )
                / 4.0
            )
            .view(B, N, N)
            .mean(dim=0)
        )

        # Calculate masked averaging
        num = (eps * mask).sum(dim=(-1, -2))
        den = mask.sum(dim=(-1, -2)).clamp(min=1)
        lddt = num / den
        if self.lddt_weight_schedule:
            t_weight = 1 + 8 * torch.relu(t - 0.5)
            lddt = (1.0 - lddt) * t_weight
            return lddt.mean()
        else:
            return (1.0 - lddt.mean()) * self.smooth_lddt_loss_weight

    def plddt_loss(
        self,
        pred_lddt,
        pred_atom_coords,
        true_atom_coords,
        true_coords_resolved_mask,
        feats,
        # multiplicity=1,
    ):
        """Compute plddt loss.
        Parameters
        ----------
        pred_lddt: torch.Tensor
            The plddt logits
        pred_atom_coords: torch.Tensor
            The predicted atom coordinates
        true_atom_coords: torch.Tensor
            The atom coordinates after symmetry correction
        true_coords_resolved_mask: torch.Tensor
            The resolved mask after symmetry correction
        feats: Dict[str, torch.Tensor]
            Dictionary containing the model input
        Returns
        -------
        torch.Tensor
            Plddt loss
        """

        # extract necessary features
        atom_mask = true_coords_resolved_mask

        R_set_to_rep_atom = feats["r_set_to_rep_atom"].float()
        # R_set_to_rep_atom = R_set_to_rep_atom.repeat_interleave(multiplicity, 0).float()

        token_type = feats["mol_type"]
        # token_type = token_type.repeat_interleave(multiplicity, 0)
        # is_nucleotide_token = (token_type == const.chain_type_ids["DNA"]).float() + (
        #     token_type == const.chain_type_ids["RNA"]
        # ).float()

        B = true_atom_coords.shape[0]

        # atom_to_token = feats["atom_to_token"].float()
        # atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)

        token_to_rep_atom = feats["token_to_rep_atom"].float()
        # token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)

        true_token_coords = torch.bmm(token_to_rep_atom, true_atom_coords)
        pred_token_coords = torch.bmm(token_to_rep_atom, pred_atom_coords)

        # compute true lddt
        true_d = torch.cdist(
            true_token_coords,
            torch.bmm(R_set_to_rep_atom, true_atom_coords),
        )
        pred_d = torch.cdist(
            pred_token_coords,
            torch.bmm(R_set_to_rep_atom, pred_atom_coords),
        )

        # compute mask
        pair_mask = atom_mask.unsqueeze(-1) * atom_mask.unsqueeze(-2)
        pair_mask = (
            pair_mask
            * (1 - torch.eye(pair_mask.shape[1], device=pair_mask.device))[None, :, :]
        )
        pair_mask = torch.einsum("bnm,bkm->bnk", pair_mask, R_set_to_rep_atom)
        pair_mask = torch.bmm(token_to_rep_atom, pair_mask)
        atom_mask = torch.bmm(token_to_rep_atom, atom_mask.unsqueeze(-1).float())
        # is_nucleotide_R_element = torch.bmm(
        #     R_set_to_rep_atom, torch.bmm(atom_to_token, is_nucleotide_token.unsqueeze(-1))
        # ).squeeze(-1)
        # cutoff = 15 + 15 * is_nucleotide_R_element.reshape(B, 1, -1).repeat(
        #     1, true_d.shape[1], 1
        # )

        # compute lddt
        target_lddt, mask_no_match = lddt_dist(
            pred_d, true_d, pair_mask, cutoff=15.0, per_atom=True
        )

        # compute loss
        num_bins = pred_lddt.shape[-1]
        bin_index = torch.floor(target_lddt * num_bins).long()
        bin_index = torch.clamp(bin_index, max=(num_bins - 1))
        lddt_one_hot = F.one_hot(bin_index, num_classes=num_bins)
        errors = -1 * torch.sum(
            lddt_one_hot * F.log_softmax(pred_lddt, dim=-1),
            dim=-1,
        )
        atom_mask = atom_mask.squeeze(-1)
        loss = torch.sum(errors * atom_mask * mask_no_match, dim=-1) / (
            1e-7 + torch.sum(atom_mask * mask_no_match, dim=-1)
        )

        # Average over the batch dimension
        loss = torch.mean(loss)

        self.log(
            "loss/plddt",
            loss.item(),
            on_epoch=True,
            logger=True,
            prog_bar=True,
            rank_zero_only=True,
        )

        return loss

    def plddt_train_step(self, batch, batch_idx):
        with torch.no_grad():
            batch = self.processor.preprocess_training(
                batch,
                esm_model=self.esm_model,
                esm_dict=self.esm_dict,
                af2_to_esm=self.af2_to_esm,
            )

            noise = torch.randn_like(batch['coords']).to(self.device)

            out_dict = self.sampler.sample(
                self.model_ema.module.forward, self.path,
                noise, batch
            )
            # out_dict = self.processor.postprocess(out_dict, batch)

            # denoised_coords = center_of_mass_norm(
            #     out_dict["denoised_coords"], batch['atom_pad_mask']
            # )
            # true_coords = center_of_mass_norm(
            #     batch['coords'], batch['atom_pad_mask']
            # )
            denoised_coords = center_random_augmentation(
                out_dict["denoised_coords"],
                batch['atom_pad_mask'],
                augmentation=False,
                centering=True,
            )
            true_coords = center_random_augmentation(
                batch['coords'],
                batch['atom_pad_mask'],
                augmentation=False,
                centering=True,
            )
            out_dict["denoised_coords"] = denoised_coords * self.processor.scale
            out_dict["coords"] = true_coords * self.processor.scale

            out_dict["true_coords_resolved_mask"] = batch["atom_resolved_mask"]

            t = torch.ones(batch['coords'].shape[0], device=self.device)
            out_feat = self.model(denoised_coords, t, batch) # use unscaled coords

        # Compute plddt loss
        plddt_out_dict = self.plddt_module(
            out_feat["latent"].detach(),
            batch,
        )

        plddt_loss = self.plddt_loss(
            plddt_out_dict["plddt_logits"],
            out_dict["denoised_coords"],
            out_dict["coords"],
            out_dict["true_coords_resolved_mask"],
            batch,
        )

        return plddt_loss

    def flow_matching_train_step(self, batch, batch_idx):
        batch = self.processor.preprocess_training(
            batch,
            esm_model=self.esm_model,
            esm_dict=self.esm_dict,
            af2_to_esm=self.af2_to_esm,
        )

        # timestep resampling
        t_size = batch['coords'].shape[0]
        t = 0.98 * logit_normal_sample(n=t_size, m=0.8, s=1.7) + 0.02 * torch.rand(t_size)
        t = t.to(self.device)
        t = t * (1 - 2 * self.t_eps) + self.t_eps

        noise = torch.randn_like(batch['coords']).to(self.device)

        _, y_t, v_t = self.path.interpolant(t, noise, batch["coords"])

        out_dict = self.model(y_t, t, batch)

        resolved_atom_mask = batch["atom_resolved_mask"].float()
        align_weights = y_t.new_ones(y_t.shape[:2])

        if self.use_rigid_align:
            with torch.no_grad(), torch.autocast("cuda", enabled=False):
                v_t = out_dict['predict_velocity'].detach().float()
                denoised_coords = y_t + v_t * (1.0 - t[:, None, None])
                coords = batch["coords"].detach().float()
                coords_aligned = weighted_rigid_align(
                    coords,
                    denoised_coords.detach().float(),
                    align_weights.detach().float(),
                    mask=resolved_atom_mask.detach().float(),
                )
                _, _, v_t_aligned = self.path.interpolant(t, noise, coords_aligned)
            target = v_t_aligned
        else:
            target = v_t

        loss = F.mse_loss(out_dict['predict_velocity'], target, reduction='none')

        loss_mask = resolved_atom_mask * align_weights
        loss = self.loss_masking(loss, loss_mask)
        loss = loss.mean()

        self.log(
            "loss/mse",
            loss.item(),
            on_epoch=True,
            logger=True,
            prog_bar=True,
            rank_zero_only=True,
        )

        if self.use_smooth_lddt_loss:
            # one-step Euler to get denoised coordinates
            denoised_coords = y_t + \
                out_dict['predict_velocity'] * (1.0 - t[:, None, None])

            # rescale coordinates to angstroms
            # denoised_coords = center_of_mass_norm(denoised_coords, batch['atom_pad_mask'])
            # true_coords = center_of_mass_norm(batch['coords'], batch['atom_pad_mask'])
            denoised_coords = center_random_augmentation(
                denoised_coords,
                batch['atom_pad_mask'],
                augmentation=False,
                centering=True,
            )
            true_coords = center_random_augmentation(
                batch['coords'],
                batch['atom_pad_mask'],
                augmentation=False,
                centering=True,
            )
            denoised_coords = denoised_coords * self.processor.scale
            true_coords = true_coords * self.processor.scale

            smooth_lddt_loss = self.smooth_lddt_loss(
                denoised_coords,
                true_coords,
                resolved_atom_mask,
                t,
            )
            loss += smooth_lddt_loss

            self.log(
                "loss/smooth_lddt",
                smooth_lddt_loss.item(),
                on_epoch=True,
                logger=True,
                prog_bar=True,
                rank_zero_only=True,
            )

        self.log(
            "loss/loss",
            loss.item(),
            on_epoch=True,
            logger=True,
            prog_bar=True,
            rank_zero_only=True,
        )

        self.log(
            "trainer/global_step",
            self.global_step,
            on_epoch=False,
            logger=True,
            prog_bar=False,
            rank_zero_only=True,
        )

        self.global_training_step = self.trainer.global_step
        self.epoch = self.trainer.current_epoch
        self.world_size = self.trainer.world_size
        return loss

    def training_step(self, batch, batch_idx):
        if self.plddt_training:
            return self.plddt_train_step(batch, batch_idx)
        else:
            return self.flow_matching_train_step(batch, batch_idx)

    @torch.no_grad()
    def validation_step(self, batch, batch_idx):
        # we skip validation step in training mode
        return

    def on_train_start(self):
        global_seed = os.environ.get("PL_GLOBAL_SEED", 42)
        pl.seed_everything(int(global_seed) + self.trainer.global_rank, True)
        return

    @torch.no_grad()
    def predict_step(self, batch, batch_idx):
        with torch.autocast(device_type='cuda', dtype=torch.float32):
            batch = self.processor.preprocess_inference(
                batch,
                esm_model=self.esm_model,
                esm_dict=self.esm_dict,
                af2_to_esm=self.af2_to_esm,
            )
            num_repeats = batch.get("num_repeats", torch.tensor(1, device=self.device)).item()
            multiplicity = batch["mol_type"].shape[0]
            num_iter = np.ceil(num_repeats / multiplicity).astype(int)
            print(f"Generating {num_repeats} samples with num_iter: {num_iter}, multiplicity: {multiplicity}")

            # num_repeats is the total number of samples to generate for one protein
            # multiplicity is the number of samples to generate at once

            curr_idx = 0
            for i in range(num_iter):
                batch_in = copy.deepcopy(batch)
                noise = torch.randn_like(batch_in['coords']).to(self.device)

                out_dict = self.sampler.sample(
                    self.model_ema.module.forward, self.path,
                    noise, batch_in
                )

                if self.plddt_module is not None:
                    denoised_coords = center_random_augmentation(
                        out_dict["denoised_coords"],
                        batch['atom_pad_mask'],
                        augmentation=False,
                        centering=True,
                    )

                    t = torch.ones(batch['coords'].shape[0], device=self.device)
                    out_feat = self.model(denoised_coords, t, batch) # use unscaled coords

                    plddt_out_dict = self.plddt_module(
                        out_feat["latent"],
                        batch_in,
                    )
                    plddts = plddt_out_dict["plddt"] * 100.0
                else:
                    plddts = None

                out_dict = self.processor.postprocess(out_dict, batch_in)

                record = Record(**batch_in['record'][0])
                gt_coord = out_dict['coords']
                sampled_coord = out_dict['denoised_coords']
                pad_mask = batch_in['atom_pad_mask']

                if num_repeats - curr_idx < multiplicity:
                    curr_num_copies = num_repeats - curr_idx
                    gt_coord = gt_coord[:curr_num_copies]
                    sampled_coord = sampled_coord[:curr_num_copies]
                    pad_mask = pad_mask[:curr_num_copies]
                    plddts = plddts[:curr_num_copies] if plddts is not None else None
                else:
                    curr_num_copies = multiplicity

                data_dir = self.trainer.datamodule.predict_dataloader().dataset.target_dir
                sample_dir = self.sample_dir
                path = Path(data_dir) / "structures" / f"{record.id}.npz"
                structure: Structure = Structure.load(path)

                for j in range(curr_num_copies):
                    file_id = j + curr_idx
                    try:
                        sampled_structure = copy.deepcopy(structure)
                        sampled_structure = process_structure(
                            sampled_structure, sampled_coord[j], pad_mask[j], record
                        )
                        # save mmcif structure
                        sampled_struct_dir = Path(sample_dir)
                        outname = f"{record.id}_sampled_{str(file_id)}"
                        save_structure(
                            sampled_structure, sampled_struct_dir, outname, 
                            plddts=plddts[j] if plddts is not None else None,
                            output_format="mmcif",
                        )
                        # save pdb structure
                        save_structure(
                            sampled_structure, sampled_struct_dir, outname, 
                            plddts=plddts[j] if plddts is not None else None,
                            output_format="pdb",
                        )
                    except:
                        print(f"Error processing {record.id}")
                        continue

                curr_idx += curr_num_copies

            self.world_size = self.trainer.world_size

        return

    def on_predict_epoch_end(self):
        dist.barrier() # wait for all processes to finish
        return

    def reset_esm(self, esm_model: str):
        self.esm_model, self.esm_dict = esm_registry[esm_model]()
        self.esm_model.eval()
        self.af2_to_esm = _af2_to_esm(self.esm_dict)
        self.esm_model = self.esm_model.to(self.device)
        self.af2_to_esm = self.af2_to_esm.to(self.device)
        print(f"Successfully reset ESM model {esm_model}")

    def setup(self, stage: str) -> None:
        """Lightning hook that is called at the beginning of fit (train + validate), validate,
        test, or predict.
        This is a good hook when you need to build models dynamically or adjust something about
        them. This hook is called on every process when using DDP.
        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
        """
        self.processor = self.hparams.processor(device=self.device)

        if self.use_esm:
            if stage == "fit" and not isinstance(
                self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
            ):
                # initialize the model with FSDP wrapper
                fsdp_params = dict(
                    mixed_precision=True,
                    flatten_parameters=True,
                    state_dict_device=torch.device("cpu"),  # reduce GPU mem usage
                    cpu_offload=True,  # enable cpu offloading
                    fp32_reduce_scatter=True,  # use fp32 reduce scatter
                )
                with enable_wrap(wrapper_cls=FSDP, **fsdp_params):
                    self.esm_model.eval()
                    # Wrap each layer in FSDP separately
                    for name, child in self.esm_model.named_children():
                        if name == "layers":
                            for layer_name, layer in child.named_children():
                                wrapped_layer = wrap(layer)
                                setattr(child, layer_name, wrapped_layer)
                    self.esm_model = wrap(self.esm_model)
                self.af2_to_esm = self.af2_to_esm.to(self.device)
            else:
                self.esm_model = self.esm_model.to(self.device)
                self.af2_to_esm = self.af2_to_esm.to(self.device)

        if stage == "fit":
            self.training_gpus = self.trainer.world_size
            self.hparams["training_gpus"] = self.training_gpus

            batch = next(iter(self.trainer.datamodule.train_dataloader()))

            batch = self.processor.preprocess_training(
                batch,
                self.esm_model,
                self.esm_dict,
                self.af2_to_esm,
            )
            y = batch["coords"]
            t = torch.zeros((y.shape[0])).cuda()

    def on_train_batch_end(self, outputs, batch, batch_idx):
        optimizer = self.optimizers()
        self.log(
            "trainer/lr",
            optimizer.param_groups[0]["lr"],
            on_epoch=True,
            logger=True,
            prog_bar=True,
            rank_zero_only=True,
        )

    def on_before_optimizer_step(self, optimizer: Optimizer) -> None:

        if isinstance(
            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
        ):

            with FullyShardedDataParallel.summon_full_params(
                self.trainer.strategy.model, with_grads=True
            ):
                clip_grad_norm_(
                    self.trainer.strategy.model.model.parameters(),
                    self.hparams.clip_grad_norm_val,
                    norm_type=2.0,
                    error_if_nonfinite=True,
                )

        else:
            clip_grad_norm_(
                self.model.parameters(),
                self.hparams.clip_grad_norm_val,
                norm_type=2.0,
                error_if_nonfinite=True,
            )
        return

    def on_before_zero_grad(self, optimizer: Optimizer) -> None:
        # if self.eval_ema:
        if isinstance(
            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
        ):

            with FullyShardedDataParallel.summon_full_params(
                self.trainer.strategy.model
            ):
                self.trainer.strategy.model.model_ema.update_parameters(
                    self.trainer.strategy.model.model
                )

        else:
            self.model_ema.update_parameters(self.model)
        return

    def on_save_checkpoint(self, checkpoint) -> None:

        if not isinstance(
            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
        ):
            layers_to_delete = []
            for k in checkpoint["state_dict"].keys():
                if k.startswith("esm_model._fsdp_wrapped_module"):
                    layers_to_delete.append(k)

            for k in layers_to_delete:
                del checkpoint["state_dict"][k]

        else:
            try:
                del checkpoint["hyper_parameters"]
            except:
                print("No hyper_parameters in checkpoint")

        return super().on_save_checkpoint(checkpoint)

    def on_load_checkpoint(self, checkpoint) -> None:

        if not isinstance(
            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
        ):
            self.training_gpus = checkpoint["hyper_parameters"]["training_gpus"]
            self.fwd_flops = checkpoint["hyper_parameters"]["fwd_flops"]

        if checkpoint["loops"] is not None:

            self.trainer.fit_loop.load_state_dict(checkpoint["loops"]["fit_loop"])
            self.trainer.validate_loop.load_state_dict(
                checkpoint["loops"]["validate_loop"]
            )
        return super().on_load_checkpoint(checkpoint)

    def configure_optimizers(self):
        optimizer = self.hparams.optimizer(
            params=self.trainer.model.parameters()
        )

        if self.hparams.scheduler is not None:
            scheduler = self.hparams.scheduler(optimizer=optimizer)
            return {
                "optimizer": optimizer,
                "lr_scheduler": {
                    "scheduler": scheduler,
                    "interval": "step",
                    "frequency": 1,
                },
            }
        return {"optimizer": optimizer}

	#
	# For licensing see accompanying LICENSE file.
	# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
	#

	import os
	import shutil
	import copy
	import numpy as np
	from pathlib import Path
	from einops import repeat

	import torch
	import torch.nn.functional as F
	from torch.optim.optimizer import Optimizer
	import torch.distributed as dist
	from torch.distributed.fsdp import FullyShardedDataParallel
	from torch.optim.swa_utils import AveragedModel
	from torch.nn.utils import clip_grad_norm_

	import lightning
	import lightning.pytorch as pl
	from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
	from fairscale.nn.wrap import enable_wrap, wrap

	from onescience.utils.simplefold.esm_utils import _af2_to_esm, esm_registry
	from onescience.datapipes.boltz_data_pipeline.types import Record, Structure
	from onescience.utils.simplefold.boltz_utils import (
	weighted_rigid_align,
	center_random_augmentation,
	process_structure,
	save_structure
	)


	def logit_normal_sample(n=1, m=0.0, s=1.0):
	# Logit-Normal Sampling from https://arxiv.org/pdf/2403.03206.pdf
	u = torch.randn(n) * s + m
	t = 1 / (1 + torch.exp(-u))
	return t


	def lddt_dist(dmat_predicted, dmat_true, mask, cutoff=15.0, per_atom=False):
	# NOTE: the mask is a pairwise mask which should have the identity elements already masked out
	# Compute mask over distances
	dists_to_score = (dmat_true < cutoff).float() * mask
	dist_l1 = torch.abs(dmat_true - dmat_predicted)

	score = 0.25 * (
	(dist_l1 < 0.5).float()
	+ (dist_l1 < 1.0).float()
	+ (dist_l1 < 2.0).float()
	+ (dist_l1 < 4.0).float()
	)

	# Normalize over the appropriate axes.
	if per_atom:
	mask_no_match = torch.sum(dists_to_score, dim=-1) != 0
	norm = 1.0 / (1e-10 + torch.sum(dists_to_score, dim=-1))
	score = norm * (1e-10 + torch.sum(dists_to_score * score, dim=-1))
	return score, mask_no_match.float()
	else:
	norm = 1.0 / (1e-10 + torch.sum(dists_to_score, dim=(-2, -1)))
	score = norm * (1e-10 + torch.sum(dists_to_score * score, dim=(-2, -1)))
	total = torch.sum(dists_to_score, dim=(-1, -2))
	return score, total


	class SimpleFold(pl.LightningModule):
	def __init__(
	self,
	architecture,
	processor,
	loss,
	path,
	sampler,
	optimizer=None,
	scheduler=None,
	plddt_module=None,
	ema_decay=0.999,
	esm_model="esm2_3B",
	aa_bolt_link=None,
	use_rigid_align=True,
	smooth_lddt_loss_weight=1.0,
	lddt_cutoff=15.0,
	clip_grad_norm_val=None,
	lddt_weight_schedule=False,
	plddt_training=False,
	sample_dir='artifacts/',
	):
	super().__init__()
	self.save_hyperparameters(logger=False)

	self.model = architecture
	self.model_ema = AveragedModel(
	self.model,
	multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(
	self.hparams.ema_decay
	),
	use_buffers=True,
	)
	self.model_ema.eval()

	self.loss = loss
	self.path = path
	self.sampler = sampler

	self.use_rigid_align = use_rigid_align
	self.lddt_cutoff = lddt_cutoff
	self.smooth_lddt_loss_weight = smooth_lddt_loss_weight
	self.use_smooth_lddt_loss = smooth_lddt_loss_weight > 0.0
	self.lddt_weight_schedule = lddt_weight_schedule
	self.plddt_training = plddt_training
	self.sample_dir = sample_dir

	self.aa_bolt_link = aa_bolt_link
	self.nval_steps = 0

	try:
	self.t_eps = self.sampler.t_eps
	except AttributeError:
	self.t_eps = 0.0

	self.use_esm = esm_model is not None
	if self.use_esm:
	self.esm_model, self.esm_dict = esm_registry[esm_model]()
	self.esm_model.eval()
	self.af2_to_esm = _af2_to_esm(self.esm_dict)
	print(f"Using ESM model: {esm_model}")
	else:
	self.esm_model = None
	self.esm_dict = None
	self.af2_to_esm = None

	self.plddt_module = plddt_module
	if self.plddt_training:
	assert self.plddt_module is not None, "PLDDT module must be provided for PLDDT training"
	self.model.eval()

	def register(self, name, tensor):
	self.register_buffer(name, tensor.type(torch.float32))

	def loss_masking(self, loss, atom_mask):
	loss_mask = repeat(atom_mask, "b s -> b s d", d=loss.shape[-1])
	loss *= loss_mask

	denom = torch.sum(atom_mask, -1, keepdim=True)
	denom = denom.unsqueeze(-1)
	loss = torch.sum(loss, dim=1, keepdim=True) / denom
	return loss

	def smooth_lddt_loss(
	self,
	pred_coords,
	true_coords,
	# is_nucleotide,
	coords_mask,
	t,
	):
	"""Compute weighted alignment.

	Parameters
	----------
	pred_coords: torch.Tensor
	The predicted atom coordinates
	true_coords: torch.Tensor
	The ground truth atom coordinates
	coords_mask: torch.Tensor
	The atoms mask

	"""
	B, N, _ = true_coords.shape
	true_dists = torch.cdist(true_coords, true_coords)

	mask = (true_dists < self.lddt_cutoff).float()
	mask = mask * (1 - torch.eye(pred_coords.shape[1], device=pred_coords.device))
	mask = mask * (coords_mask.unsqueeze(-1) * coords_mask.unsqueeze(-2))

	# Compute distances between all pairs of atoms
	pred_dists = torch.cdist(pred_coords, pred_coords)
	dist_diff = torch.abs(true_dists - pred_dists)

	# Compute epsilon values
	eps = (
	(
	(
	F.sigmoid(0.5 - dist_diff)
	+ F.sigmoid(1.0 - dist_diff)
	+ F.sigmoid(2.0 - dist_diff)
	+ F.sigmoid(4.0 - dist_diff)
	)
	/ 4.0
	)
	.view(B, N, N)
	.mean(dim=0)
	)

	# Calculate masked averaging
	num = (eps * mask).sum(dim=(-1, -2))
	den = mask.sum(dim=(-1, -2)).clamp(min=1)
	lddt = num / den
	if self.lddt_weight_schedule:
	t_weight = 1 + 8 * torch.relu(t - 0.5)
	lddt = (1.0 - lddt) * t_weight
	return lddt.mean()
	else:
	return (1.0 - lddt.mean()) * self.smooth_lddt_loss_weight

	def plddt_loss(
	self,
	pred_lddt,
	pred_atom_coords,
	true_atom_coords,
	true_coords_resolved_mask,
	feats,
	# multiplicity=1,
	):
	"""Compute plddt loss.

	Parameters
	----------
	pred_lddt: torch.Tensor
	The plddt logits
	pred_atom_coords: torch.Tensor
	The predicted atom coordinates
	true_atom_coords: torch.Tensor
	The atom coordinates after symmetry correction
	true_coords_resolved_mask: torch.Tensor
	The resolved mask after symmetry correction
	feats: Dict[str, torch.Tensor]
	Dictionary containing the model input

	Returns
	-------
	torch.Tensor
	Plddt loss

	"""

	# extract necessary features
	atom_mask = true_coords_resolved_mask

	R_set_to_rep_atom = feats["r_set_to_rep_atom"].float()
	# R_set_to_rep_atom = R_set_to_rep_atom.repeat_interleave(multiplicity, 0).float()

	token_type = feats["mol_type"]
	# token_type = token_type.repeat_interleave(multiplicity, 0)
	# is_nucleotide_token = (token_type == const.chain_type_ids["DNA"]).float() + (
	# token_type == const.chain_type_ids["RNA"]
	# ).float()

	B = true_atom_coords.shape[0]

	# atom_to_token = feats["atom_to_token"].float()
	# atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)

	token_to_rep_atom = feats["token_to_rep_atom"].float()
	# token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)

	true_token_coords = torch.bmm(token_to_rep_atom, true_atom_coords)
	pred_token_coords = torch.bmm(token_to_rep_atom, pred_atom_coords)

	# compute true lddt
	true_d = torch.cdist(
	true_token_coords,
	torch.bmm(R_set_to_rep_atom, true_atom_coords),
	)
	pred_d = torch.cdist(
	pred_token_coords,
	torch.bmm(R_set_to_rep_atom, pred_atom_coords),
	)

	# compute mask
	pair_mask = atom_mask.unsqueeze(-1) * atom_mask.unsqueeze(-2)
	pair_mask = (
	pair_mask
	* (1 - torch.eye(pair_mask.shape[1], device=pair_mask.device))[None, :, :]
	)
	pair_mask = torch.einsum("bnm,bkm->bnk", pair_mask, R_set_to_rep_atom)
	pair_mask = torch.bmm(token_to_rep_atom, pair_mask)
	atom_mask = torch.bmm(token_to_rep_atom, atom_mask.unsqueeze(-1).float())
	# is_nucleotide_R_element = torch.bmm(
	# R_set_to_rep_atom, torch.bmm(atom_to_token, is_nucleotide_token.unsqueeze(-1))
	# ).squeeze(-1)
	# cutoff = 15 + 15 * is_nucleotide_R_element.reshape(B, 1, -1).repeat(
	# 1, true_d.shape[1], 1
	# )

	# compute lddt
	target_lddt, mask_no_match = lddt_dist(
	pred_d, true_d, pair_mask, cutoff=15.0, per_atom=True
	)

	# compute loss
	num_bins = pred_lddt.shape[-1]
	bin_index = torch.floor(target_lddt * num_bins).long()
	bin_index = torch.clamp(bin_index, max=(num_bins - 1))
	lddt_one_hot = F.one_hot(bin_index, num_classes=num_bins)
	errors = -1 * torch.sum(
	lddt_one_hot * F.log_softmax(pred_lddt, dim=-1),
	dim=-1,
	)
	atom_mask = atom_mask.squeeze(-1)
	loss = torch.sum(errors * atom_mask * mask_no_match, dim=-1) / (
	1e-7 + torch.sum(atom_mask * mask_no_match, dim=-1)
	)

	# Average over the batch dimension
	loss = torch.mean(loss)

	self.log(
	"loss/plddt",
	loss.item(),
	on_epoch=True,
	logger=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	return loss

	def plddt_train_step(self, batch, batch_idx):
	with torch.no_grad():
	batch = self.processor.preprocess_training(
	batch,
	esm_model=self.esm_model,
	esm_dict=self.esm_dict,
	af2_to_esm=self.af2_to_esm,
	)

	noise = torch.randn_like(batch['coords']).to(self.device)

	out_dict = self.sampler.sample(
	self.model_ema.module.forward, self.path,
	noise, batch
	)
	# out_dict = self.processor.postprocess(out_dict, batch)

	# denoised_coords = center_of_mass_norm(
	# out_dict["denoised_coords"], batch['atom_pad_mask']
	# )
	# true_coords = center_of_mass_norm(
	# batch['coords'], batch['atom_pad_mask']
	# )
	denoised_coords = center_random_augmentation(
	out_dict["denoised_coords"],
	batch['atom_pad_mask'],
	augmentation=False,
	centering=True,
	)
	true_coords = center_random_augmentation(
	batch['coords'],
	batch['atom_pad_mask'],
	augmentation=False,
	centering=True,
	)
	out_dict["denoised_coords"] = denoised_coords * self.processor.scale
	out_dict["coords"] = true_coords * self.processor.scale

	out_dict["true_coords_resolved_mask"] = batch["atom_resolved_mask"]

	t = torch.ones(batch['coords'].shape[0], device=self.device)
	out_feat = self.model(denoised_coords, t, batch) # use unscaled coords

	# Compute plddt loss
	plddt_out_dict = self.plddt_module(
	out_feat["latent"].detach(),
	batch,
	)

	plddt_loss = self.plddt_loss(
	plddt_out_dict["plddt_logits"],
	out_dict["denoised_coords"],
	out_dict["coords"],
	out_dict["true_coords_resolved_mask"],
	batch,
	)

	return plddt_loss

	def flow_matching_train_step(self, batch, batch_idx):
	batch = self.processor.preprocess_training(
	batch,
	esm_model=self.esm_model,
	esm_dict=self.esm_dict,
	af2_to_esm=self.af2_to_esm,
	)

	# timestep resampling
	t_size = batch['coords'].shape[0]
	t = 0.98 * logit_normal_sample(n=t_size, m=0.8, s=1.7) + 0.02 * torch.rand(t_size)
	t = t.to(self.device)
	t = t * (1 - 2 * self.t_eps) + self.t_eps

	noise = torch.randn_like(batch['coords']).to(self.device)

	_, y_t, v_t = self.path.interpolant(t, noise, batch["coords"])

	out_dict = self.model(y_t, t, batch)

	resolved_atom_mask = batch["atom_resolved_mask"].float()
	align_weights = y_t.new_ones(y_t.shape[:2])

	if self.use_rigid_align:
	with torch.no_grad(), torch.autocast("cuda", enabled=False):
	v_t = out_dict['predict_velocity'].detach().float()
	denoised_coords = y_t + v_t * (1.0 - t[:, None, None])
	coords = batch["coords"].detach().float()
	coords_aligned = weighted_rigid_align(
	coords,
	denoised_coords.detach().float(),
	align_weights.detach().float(),
	mask=resolved_atom_mask.detach().float(),
	)
	_, _, v_t_aligned = self.path.interpolant(t, noise, coords_aligned)
	target = v_t_aligned
	else:
	target = v_t

	loss = F.mse_loss(out_dict['predict_velocity'], target, reduction='none')

	loss_mask = resolved_atom_mask * align_weights
	loss = self.loss_masking(loss, loss_mask)
	loss = loss.mean()

	self.log(
	"loss/mse",
	loss.item(),
	on_epoch=True,
	logger=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	if self.use_smooth_lddt_loss:
	# one-step Euler to get denoised coordinates
	denoised_coords = y_t + \
	out_dict['predict_velocity'] * (1.0 - t[:, None, None])

	# rescale coordinates to angstroms
	# denoised_coords = center_of_mass_norm(denoised_coords, batch['atom_pad_mask'])
	# true_coords = center_of_mass_norm(batch['coords'], batch['atom_pad_mask'])
	denoised_coords = center_random_augmentation(
	denoised_coords,
	batch['atom_pad_mask'],
	augmentation=False,
	centering=True,
	)
	true_coords = center_random_augmentation(
	batch['coords'],
	batch['atom_pad_mask'],
	augmentation=False,
	centering=True,
	)
	denoised_coords = denoised_coords * self.processor.scale
	true_coords = true_coords * self.processor.scale

	smooth_lddt_loss = self.smooth_lddt_loss(
	denoised_coords,
	true_coords,
	resolved_atom_mask,
	t,
	)
	loss += smooth_lddt_loss

	self.log(
	"loss/smooth_lddt",
	smooth_lddt_loss.item(),
	on_epoch=True,
	logger=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	self.log(
	"loss/loss",
	loss.item(),
	on_epoch=True,
	logger=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	self.log(
	"trainer/global_step",
	self.global_step,
	on_epoch=False,
	logger=True,
	prog_bar=False,
	rank_zero_only=True,
	)

	self.global_training_step = self.trainer.global_step
	self.epoch = self.trainer.current_epoch
	self.world_size = self.trainer.world_size
	return loss

	def training_step(self, batch, batch_idx):
	if self.plddt_training:
	return self.plddt_train_step(batch, batch_idx)
	else:
	return self.flow_matching_train_step(batch, batch_idx)

	@torch.no_grad()
	def validation_step(self, batch, batch_idx):
	# we skip validation step in training mode
	return

	def on_train_start(self):
	global_seed = os.environ.get("PL_GLOBAL_SEED", 42)
	pl.seed_everything(int(global_seed) + self.trainer.global_rank, True)
	return

	@torch.no_grad()
	def predict_step(self, batch, batch_idx):
	with torch.autocast(device_type='cuda', dtype=torch.float32):
	batch = self.processor.preprocess_inference(
	batch,
	esm_model=self.esm_model,
	esm_dict=self.esm_dict,
	af2_to_esm=self.af2_to_esm,
	)
	num_repeats = batch.get("num_repeats", torch.tensor(1, device=self.device)).item()
	multiplicity = batch["mol_type"].shape[0]
	num_iter = np.ceil(num_repeats / multiplicity).astype(int)
	print(f"Generating {num_repeats} samples with num_iter: {num_iter}, multiplicity: {multiplicity}")

	# num_repeats is the total number of samples to generate for one protein
	# multiplicity is the number of samples to generate at once

	curr_idx = 0
	for i in range(num_iter):
	batch_in = copy.deepcopy(batch)
	noise = torch.randn_like(batch_in['coords']).to(self.device)

	out_dict = self.sampler.sample(
	self.model_ema.module.forward, self.path,
	noise, batch_in
	)

	if self.plddt_module is not None:
	denoised_coords = center_random_augmentation(
	out_dict["denoised_coords"],
	batch['atom_pad_mask'],
	augmentation=False,
	centering=True,
	)

	t = torch.ones(batch['coords'].shape[0], device=self.device)
	out_feat = self.model(denoised_coords, t, batch) # use unscaled coords

	plddt_out_dict = self.plddt_module(
	out_feat["latent"],
	batch_in,
	)
	plddts = plddt_out_dict["plddt"] * 100.0
	else:
	plddts = None

	out_dict = self.processor.postprocess(out_dict, batch_in)

	record = Record(**batch_in['record'][0])
	gt_coord = out_dict['coords']
	sampled_coord = out_dict['denoised_coords']
	pad_mask = batch_in['atom_pad_mask']

	if num_repeats - curr_idx < multiplicity:
	curr_num_copies = num_repeats - curr_idx
	gt_coord = gt_coord[:curr_num_copies]
	sampled_coord = sampled_coord[:curr_num_copies]
	pad_mask = pad_mask[:curr_num_copies]
	plddts = plddts[:curr_num_copies] if plddts is not None else None
	else:
	curr_num_copies = multiplicity

	data_dir = self.trainer.datamodule.predict_dataloader().dataset.target_dir
	sample_dir = self.sample_dir
	path = Path(data_dir) / "structures" / f"{record.id}.npz"
	structure: Structure = Structure.load(path)

	for j in range(curr_num_copies):
	file_id = j + curr_idx
	try:
	sampled_structure = copy.deepcopy(structure)
	sampled_structure = process_structure(
	sampled_structure, sampled_coord[j], pad_mask[j], record
	)
	# save mmcif structure
	sampled_struct_dir = Path(sample_dir)
	outname = f"{record.id}_sampled_{str(file_id)}"
	save_structure(
	sampled_structure, sampled_struct_dir, outname,
	plddts=plddts[j] if plddts is not None else None,
	output_format="mmcif",
	)
	# save pdb structure
	save_structure(
	sampled_structure, sampled_struct_dir, outname,
	plddts=plddts[j] if plddts is not None else None,
	output_format="pdb",
	)
	except:
	print(f"Error processing {record.id}")
	continue

	curr_idx += curr_num_copies

	self.world_size = self.trainer.world_size

	return

	def on_predict_epoch_end(self):
	dist.barrier() # wait for all processes to finish
	return

	def reset_esm(self, esm_model: str):
	self.esm_model, self.esm_dict = esm_registry[esm_model]()
	self.esm_model.eval()
	self.af2_to_esm = _af2_to_esm(self.esm_dict)
	self.esm_model = self.esm_model.to(self.device)
	self.af2_to_esm = self.af2_to_esm.to(self.device)
	print(f"Successfully reset ESM model {esm_model}")

	def setup(self, stage: str) -> None:
	"""Lightning hook that is called at the beginning of fit (train + validate), validate,
	test, or predict.

	This is a good hook when you need to build models dynamically or adjust something about
	them. This hook is called on every process when using DDP.

	:param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
	"""
	self.processor = self.hparams.processor(device=self.device)

	if self.use_esm:
	if stage == "fit" and not isinstance(
	self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
	):
	# initialize the model with FSDP wrapper
	fsdp_params = dict(
	mixed_precision=True,
	flatten_parameters=True,
	state_dict_device=torch.device("cpu"), # reduce GPU mem usage
	cpu_offload=True, # enable cpu offloading
	fp32_reduce_scatter=True, # use fp32 reduce scatter
	)
	with enable_wrap(wrapper_cls=FSDP, **fsdp_params):
	self.esm_model.eval()
	# Wrap each layer in FSDP separately
	for name, child in self.esm_model.named_children():
	if name == "layers":
	for layer_name, layer in child.named_children():
	wrapped_layer = wrap(layer)
	setattr(child, layer_name, wrapped_layer)
	self.esm_model = wrap(self.esm_model)
	self.af2_to_esm = self.af2_to_esm.to(self.device)
	else:
	self.esm_model = self.esm_model.to(self.device)
	self.af2_to_esm = self.af2_to_esm.to(self.device)

	if stage == "fit":
	self.training_gpus = self.trainer.world_size
	self.hparams["training_gpus"] = self.training_gpus

	batch = next(iter(self.trainer.datamodule.train_dataloader()))

	batch = self.processor.preprocess_training(
	batch,
	self.esm_model,
	self.esm_dict,
	self.af2_to_esm,
	)
	y = batch["coords"]
	t = torch.zeros((y.shape[0])).cuda()

	def on_train_batch_end(self, outputs, batch, batch_idx):
	optimizer = self.optimizers()
	self.log(
	"trainer/lr",
	optimizer.param_groups[0]["lr"],
	on_epoch=True,
	logger=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	def on_before_optimizer_step(self, optimizer: Optimizer) -> None:

	if isinstance(
	self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
	):

	with FullyShardedDataParallel.summon_full_params(
	self.trainer.strategy.model, with_grads=True
	):
	clip_grad_norm_(
	self.trainer.strategy.model.model.parameters(),
	self.hparams.clip_grad_norm_val,
	norm_type=2.0,
	error_if_nonfinite=True,
	)

	else:
	clip_grad_norm_(
	self.model.parameters(),
	self.hparams.clip_grad_norm_val,
	norm_type=2.0,
	error_if_nonfinite=True,
	)
	return

	def on_before_zero_grad(self, optimizer: Optimizer) -> None:
	# if self.eval_ema:
	if isinstance(
	self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
	):

	with FullyShardedDataParallel.summon_full_params(
	self.trainer.strategy.model
	):
	self.trainer.strategy.model.model_ema.update_parameters(
	self.trainer.strategy.model.model
	)

	else:
	self.model_ema.update_parameters(self.model)
	return

	def on_save_checkpoint(self, checkpoint) -> None:

	if not isinstance(
	self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
	):
	layers_to_delete = []
	for k in checkpoint["state_dict"].keys():
	if k.startswith("esm_model._fsdp_wrapped_module"):
	layers_to_delete.append(k)

	for k in layers_to_delete:
	del checkpoint["state_dict"][k]

	else:
	try:
	del checkpoint["hyper_parameters"]
	except:
	print("No hyper_parameters in checkpoint")

	return super().on_save_checkpoint(checkpoint)

	def on_load_checkpoint(self, checkpoint) -> None:

	if not isinstance(
	self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
	):
	self.training_gpus = checkpoint["hyper_parameters"]["training_gpus"]
	self.fwd_flops = checkpoint["hyper_parameters"]["fwd_flops"]

	if checkpoint["loops"] is not None:

	self.trainer.fit_loop.load_state_dict(checkpoint["loops"]["fit_loop"])
	self.trainer.validate_loop.load_state_dict(
	checkpoint["loops"]["validate_loop"]
	)
	return super().on_load_checkpoint(checkpoint)

	def configure_optimizers(self):
	optimizer = self.hparams.optimizer(
	params=self.trainer.model.parameters()
	)

	if self.hparams.scheduler is not None:
	scheduler = self.hparams.scheduler(optimizer=optimizer)
	return {
	"optimizer": optimizer,
	"lr_scheduler": {
	"scheduler": scheduler,
	"interval": "step",
	"frequency": 1,
	},
	}
	return {"optimizer": optimizer}