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