//  [python]
def forward(
        self,
        feature_dict: Dict[str, torch.Tensor],
        nloops: int = 5,
        refine_steps: int = 0,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        x = feature_dict["msa"]
        x2 = feature_dict["cov"]
        nres = x.size()[1]

        x = self.embed(x)
        x = self.vgru(x)[0]
        x = self.hgru(x[-1, :, :].unsqueeze(1))[0]
        mat1d = x.permute(1, 2, 0)
        x = mat1d.unsqueeze(2) * mat1d.unsqueeze(3)

        resinp = torch.cat((x, x2), dim=1)

        if x.requires_grad:
            # now call the checkpoint API and get the output
            x = checkpoint_sequential(self.resnet, 4, resinp)
        else:
            x = self.resnet(resinp)

        dm = x[:, 0, :, :]
        conf = x[:, 1, :, :].mean(dim=2)
        best_conf = conf

        # Ensure symmetric matrix
        dm = (dm + dm.transpose(1, 2)) / 2
        # Force values to be non-negative
        dm = torch.abs(dm)
        # See https://math.stackexchange.com/questions/156161/finding-the-coordinates-of-points-from-distance-matrix
        M = 0.5 * (
            dm[:, 0:1, :].expand(-1, nres, -1) ** 2
            + dm[:, :, 0:1].expand(-1, -1, nres) ** 2
            - dm**2
        )
        w, v = torch.linalg.eigh(M.float())
        w = torch.clamp(F.relu(w, inplace=False), min=1e-8)
        w = torch.diag_embed(w.sqrt())
        mds_coords = torch.matmul(v, w)[:, :, -8:]
        coordembed = torch.cat((mat1d.permute(0, 2, 1), mds_coords), dim=2)

        gru_out = self.coord_gru(coordembed)[0]

        ca_coords = self.coord_fc(gru_out)

        if refine_steps > 0:
            ca_coords = refine_coords(ca_coords.squeeze(0), refine_steps).unsqueeze(0)

        best_coords = ca_coords

        # print(best_conf.sum().item())

        for i in range(nloops):
            # print(ca_coords.size())

            # Add small amount of noise to seed coordinates
            # ca_coords = ca_coords + 0.5 * torch.randn_like(ca_coords)
            # if refine_steps > 0:
            #    ca_coords = refine_coords(ca_coords.squeeze(0), refine_steps).unsqueeze(0)
            dmap = torch.clamp(
                (ca_coords - ca_coords.transpose(0, 1)).pow(2).sum(dim=2), min=1e-8
            ).sqrt()
            resinp = torch.cat((resinp[:, :-1], dmap.unsqueeze(0).unsqueeze(0)), dim=1)

            if resinp.requires_grad:
                # now call the checkpoint API and get the output
                x = checkpoint_sequential(self.resnet, 4, resinp)
            else:
                x = self.resnet(resinp)

            dm = x[:, 0, :, :]
            conf = x[:, 1, :, :].mean(dim=2)

            # print(conf.sum().item())

            # Ensure symmetric matrix
            dm = (dm + dm.transpose(1, 2)) / 2
            # Force values to be non-negative
            dm = torch.abs(dm)
            # See https://math.stackexchange.com/questions/156161/finding-the-coordinates-of-points-from-distance-matrix
            M = 0.5 * (
                dm[:, 0:1, :].expand(-1, nres, -1) ** 2
                + dm[:, :, 0:1].expand(-1, -1, nres) ** 2
                - dm**2
            )
            w, v = torch.linalg.eigh(M.float())
            w = torch.clamp(F.relu(w, inplace=False), min=1e-8)
            w = torch.diag_embed(w.sqrt())
            mds_coords = torch.matmul(v, w)[:, :, -8:]
            coordembed = torch.cat((mat1d.permute(0, 2, 1), mds_coords), dim=2)

            gru_out = self.coord_gru(coordembed)[0]

            ca_coords = self.coord_fc(gru_out)

            if conf.mean() > best_conf.mean():
                # if refine_steps > 0:
                #    ca_coords = refine_coords(ca_coords.squeeze(0), refine_steps).unsqueeze(0)
                best_conf = conf
                best_coords = ca_coords

        if refine_steps > 0:
            best_coords = refine_coords(best_coords.squeeze(0), refine_steps).unsqueeze(
                0
            )

        coords_out = calpha_to_main_chain(best_coords)
        conf_out = torch.sigmoid(best_conf)

        return coords_out, conf_out