//  [python]
def forward(self, x, x2, nloops=5, refine_steps=0):

        nseqs = x.size()[0]
        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