given the code
from spin_glass_base import SpinGlassBase, DistributionType
import numpy as np
from typing import Dict, Any, Optional

class TwoDSpinGlass(SpinGlassBase):
    def __init__(self, N: int, distribution: DistributionType = DistributionType.BIMODAL,
                 spins: Optional[np.ndarray] = None, horizontal_couplings: Optional[np.ndarray] = None,
                 vertical_couplings: Optional[np.ndarray] = None):
        self.horizontal_couplings = horizontal_couplings
        self.vertical_couplings = vertical_couplings
        super().__init__(N, distribution, spins)

    def initialize_spins(self) -> np.ndarray:
        return np.random.choice([-1, 1], size=(self.N, self.N))

    def generate_couplings(self) -> Dict[str, np.ndarray]:
        if self.horizontal_couplings is None:
            self.horizontal_couplings = self._generate_distribution((self.N, self.N))
        if self.vertical_couplings is None:
            self.vertical_couplings = self._generate_distribution((self.N, self.N))
        return {'horizontal': self.horizontal_couplings, 'vertical': self.vertical_couplings}

    def calculate_energy(self) -> float:
        horizontal_energy = np.sum(self.horizontal_couplings * self.spins * np.roll(self.spins, shift=-1, axis=1))
        vertical_energy = np.sum(self.vertical_couplings * self.spins * np.roll(self.spins, shift=-1, axis=0))
        return -(horizontal_energy + vertical_energy)

    def to_dict(self) -> Dict[str, Any]:
        return {
            'N': self.N,
            'id': self.id,
            'problem_type': '2D Spin-Glass',
            'distribution': self.distribution.value,
            'spins': self.spins.tolist(),
            'horizontal_couplings': self.horizontal_couplings.tolist(),
            'vertical_couplings': self.vertical_couplings.tolist()
        }

i've built the following digital annealer using the gpu from the paper : Physics-Inspired Optimization for Quadratic Unconstrained Problems Using a Digital Annealer
import time
import torch
from SpinGlassSolution import SpinGlassSolution
from SpinGlassSolver import SpinGlassSolver
from spin_glass_base import SpinGlassBase
from sk_spin_glass import SKSpinGlass
from two_d_spin_glass import TwoDSpinGlass

class DigitalAnnealerSpinGlassSolver(SpinGlassSolver):
    def solve(self, instance: SpinGlassBase) -> SpinGlassSolution:
        if not isinstance(instance, (SKSpinGlass, TwoDSpinGlass)):
            raise ValueError("This solver only supports SK and 2D Lattice Spin Glass instances")

        start_time = time.time()
        initial_temperature = 100.0
        final_temperature = 0.01
        alpha = 0.95
        num_iterations = 250
        num_runs = 500
        offset_increase_rate = 1
        N = instance.N
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Prepare coupling matrices
        if isinstance(instance, SKSpinGlass):
            J = torch.tensor(instance.couplings, dtype=torch.float32, device=device)
        else:  # TwoDSpinGlass
            J_h = torch.tensor(instance.horizontal_couplings, dtype=torch.float32, device=device)
            J_v = torch.tensor(instance.vertical_couplings, dtype=torch.float32, device=device)

        best_energy = float('inf')
        best_spins = None

        # Keep initial state on CPU
        initial_state = torch.tensor(instance.spins, dtype=torch.float32)

        # Main loop for multiple runs
        for _ in range(num_runs):
            spins = initial_state.clone()
            temperature = initial_temperature
            Eoffset = 0.0

            # Monte Carlo steps
            for step in range(num_iterations):
                if temperature > final_temperature:
                    temperature *= alpha

                # Calculate energy changes ΔE for flipping each spin
                # spins_gpu = spins.to(device)
                # if isinstance(instance, SKSpinGlass):
                #     delta_Es = 2 * spins_gpu * torch.matmul(J, spins_gpu)
                # else:  # TwoDSpinGlass
                #     spins_2d = spins_gpu.view(N, N)
                #     delta_Es = (2 * spins_2d * (
                #         torch.matmul(J_h, spins_2d.roll(1, dims=1)) +
                #         torch.matmul(spins_2d.roll(1, dims=0), J_v)
                #     )).view(-1)
                # delta_Es = delta_Es.cpu()
                
                if isinstance(instance, SKSpinGlass):
                    spins_gpu = spins.to(device)  # Keep spins as 1D
                    delta_Es = 2 * spins_gpu * torch.matmul(J, spins_gpu)
                else:  # TwoDSpinGlass
                    spins_2d = spins_gpu.view(N, N)  # Keep spins as 2D array
                    delta_Es = 2 * spins_2d * (
                        torch.matmul(J_h, spins_2d.roll(1, dims=1)) +
                        torch.matmul(spins_2d.roll(1, dims=0))
                    )
                delta_Es = delta_Es.view(-1).cpu()  # Move result back to CPU if needed

                # Adjust ΔE with the offset and decide on flips (on CPU)
                accept_probabilities = torch.exp(-(delta_Es - Eoffset) / temperature)
                accepted_flips = (delta_Es < Eoffset) | (torch.rand(N * N if isinstance(instance, TwoDSpinGlass) else N) < accept_probabilities)

                if torch.any(accepted_flips):
                    accepted_indices = torch.nonzero(accepted_flips).squeeze()
                    if accepted_indices.dim() == 0:
                        selected_flip = accepted_indices.item()
                    else:
                        accepted_flips = accepted_flips.view(N, N)  # Keep flips as 2D
                        selected_flip = torch.nonzero(accepted_flips).squeeze()
                        #selected_flip = accepted_indices[torch.randint(len(accepted_indices), (1,)).item()]
                        if selected_flip.dim() == 1:
                            spins_2d[selected_flip[0], selected_flip[1]] *= -1  # Flip 2D spins
                    spins[selected_flip] *= -1
                    Eoffset = 0.0
                else:
                    Eoffset += offset_increase_rate

            # Calculate the final energy for this run
            spins_gpu = spins.to(device)
            if isinstance(instance, SKSpinGlass):
                final_energy = -0.5 * torch.sum(spins_gpu * torch.matmul(J, spins_gpu)).item()
            else:  # TwoDSpinGlass
                spins_2d = spins_gpu.view(N, N)
                final_energy = -(
                    torch.sum(J_h * spins_2d * spins_2d.roll(1, dims=1)) +
                    torch.sum(J_v * spins_2d * spins_2d.roll(1, dims=0))
                ).item()

            if final_energy < best_energy:
                best_energy = final_energy
                best_spins = spins.clone()

        solving_time = time.time() - start_time

        # Return final result
        final_spins = best_spins.numpy()
        if isinstance(instance, TwoDSpinGlass):
            final_spins = final_spins.reshape(N, N)
        return SpinGlassSolution(final_spins, best_energy, solving_time)

this solution has a problem with the two-d sping glass while it works good for the SK 
from spin_glass_base import SpinGlassBase, DistributionType
import numpy as np
from typing import Dict, Any, Optional

class SKSpinGlass(SpinGlassBase):
    def __init__(self, N: int, distribution: DistributionType = DistributionType.BIMODAL,
                 spins: Optional[np.ndarray] = None, couplings: Optional[np.ndarray] = None):
        self.couplings = couplings
        super().__init__(N, distribution, spins)

    def initialize_spins(self) -> np.ndarray:
        return np.random.choice([-1, 1], size=self.N)

    def generate_couplings(self) -> np.ndarray:
        if self.couplings is not None:
            return self.couplings

        couplings = self._generate_distribution((self.N, self.N))
        couplings = np.triu(couplings, 1)
        couplings += couplings.T
        np.fill_diagonal(couplings, 0)
        return couplings

    def calculate_energy(self) -> float:
        return -np.sum(np.triu(self.couplings, 1) * np.outer(self.spins, self.spins))

    def to_dict(self) -> Dict[str, Any]:
        return {
            'N': self.N,
            'id': self.id,
            'problem_type': 'SK Spin-Glass',
            'distribution': self.distribution.value,
            'spins': self.spins.tolist(),
            'couplings': self.couplings.tolist()
        }

the main reason is because the two-d problem has spins as a 2 array while sk a 1d array