diff --git a/ivon/_ivon.py b/ivon/_ivon.py
index 5dfb6cf..2b57474 100644
--- a/ivon/_ivon.py
+++ b/ivon/_ivon.py
@@ -11,7 +11,17 @@
 
 
 def _welford_mean(avg: Optional[Tensor], newval: Tensor, count: int) -> Tensor:
-    return newval if avg is None else avg + (newval - avg) / count
+    # Use autocast for numerical stability during arithmetic operations
+    with torch.autocast(device_type=newval.device.type):
+        if avg is None:
+            return newval
+        else:
+            # Ensure consistent device and dtype
+            if avg.device != newval.device:
+                avg = avg.to(device=newval.device)
+            if avg.dtype != newval.dtype:
+                newval = newval.to(dtype=avg.dtype)
+            return avg + (newval - avg) / count
 
 
 class IVON(torch.optim.Optimizer):
@@ -74,7 +84,7 @@ def __init__(
         self.mc_samples = mc_samples
         self.hess_approx = hess_approx
         self.sync = sync
-        self._numel, self._device, self._dtype = self._get_param_configs()
+        self._numel = self._get_param_configs()
         self.current_step = 0
         self.debias = debias
         self.rescale_lr = rescale_lr
@@ -85,45 +95,55 @@ def __init__(
         self._init_buffers()
 
     def _get_param_configs(self):
-        all_params = []
         for pg in self.param_groups:
             pg["numel"] = sum(p.numel() for p in pg["params"] if p is not None)
-            all_params += [p for p in pg["params"] if p is not None]
-        if len(all_params) == 0:
-            return 0, torch.device("cpu"), torch.get_default_dtype()
-        devices = {p.device for p in all_params}
-        if len(devices) > 1:
-            raise ValueError(
-                "Parameters are on different devices: "
-                f"{[str(d) for d in devices]}"
-            )
-        device = next(iter(devices))
-        dtypes = {p.dtype for p in all_params}
-        if len(dtypes) > 1:
-            raise ValueError(
-                "Parameters are on different dtypes: "
-                f"{[str(d) for d in dtypes]}"
-            )
-        dtype = next(iter(dtypes))
         total = sum(pg["numel"] for pg in self.param_groups)
-        return total, device, dtype
+        return total
 
-    def _reset_samples(self):
-        self.state['count'] = 0
-        self.state['avg_grad'] = None
-        self.state['avg_nxg'] = None
-        self.state['avg_gsq'] = None
+    @property
+    def _device(self):
+        """Get device from first parameter"""
+        for pg in self.param_groups:
+            for p in pg["params"]:
+                if p is not None:
+                    return p.device
+        return torch.device("cpu")
+    
+    @property
+    def _dtype(self):
+        """Get dtype from first parameter"""
+        for pg in self.param_groups:
+            for p in pg["params"]:
+                if p is not None:
+                    return p.dtype
+        return torch.get_default_dtype()
+
+    def _reset_samples(self, reset_count=True):
+        # Do not reset critical state unless explicitly told to do so
+        if reset_count:
+            self.state['count'] = 0
+        # Intentionally keep gradient-related states as they are
 
     def _init_buffers(self):
         for group in self.param_groups:
             hess_init, numel = group["hess_init"], group["numel"]
 
+            # Always use float32 for internal optimizer buffers for numerical stability
             group["momentum"] = torch.zeros(
-                numel, device=self._device, dtype=self._dtype
+                numel, device=self._device, dtype=torch.float32
             )
             group["hess"] = torch.zeros(
-                numel, device=self._device, dtype=self._dtype
+                numel, device=self._device, dtype=torch.float32
             ).add(torch.as_tensor(hess_init))
+    
+    def _ensure_buffers_on_device(self):
+        """Ensure buffers are on the same device as parameters"""
+        for group in self.param_groups:
+            # Always use float32 for internal optimizer buffers
+            if "momentum" in group and group["momentum"].device != self._device:
+                group["momentum"] = group["momentum"].to(device=self._device, dtype=torch.float32)
+            if "hess" in group and group["hess"].device != self._device:
+                group["hess"] = group["hess"].to(device=self._device, dtype=torch.float32)
 
     @contextmanager
     def sampled_params(self, train: bool = False):
@@ -143,19 +163,53 @@ def _restore_param_average(
 
                 p_slice = slice(offset, offset + p.numel())
 
-                p.data = param_avg[p_slice].view(p.shape)
+                # Convert param_avg back to parameter's original dtype
+                p.data = param_avg[p_slice].view(p.shape).to(dtype=p.dtype)
                 if train:
-                    if p.requires_grad:
+                    # Always collect gradient, handling potential accumulated gradients
+                    if p.grad is not None:
                         param_grads.append(p.grad.flatten())
                     else:
-                        param_grads.append(torch.zeros_like(p).flatten())
+                        param_grads.append(torch.zeros(
+                            p.numel(), 
+                            device=p.device, 
+                            dtype=p.dtype
+                        ))
                 offset += p.numel()
         assert offset == self._numel  # sanity check
 
-        if train:  # collect grad sample for training
+        if train and param_grads:  # collect grad sample for training
             grad_sample = torch.cat(param_grads, 0)
+            
+            # Ensure state is initialized
+            if 'count' not in self.state:
+                self.state['count'] = 0
+            
+            # Increment count
             count = self.state["count"] + 1
             self.state["count"] = count
+            
+            # Compute total number of parameters
+            total_numel = sum(group['numel'] for group in self.param_groups)
+            
+            # Initialize gradient-related states if not already done
+            if self.state.get('avg_grad') is None:
+                # Use gradient dtype for state initialization
+                self.state['avg_grad'] = torch.zeros(
+                    total_numel, 
+                    device=self._device, 
+                    dtype=grad_sample.dtype
+                )
+                self.state['avg_nxg'] = torch.zeros_like(self.state['avg_grad'])
+                self.state['avg_gsq'] = torch.zeros_like(self.state['avg_grad'])
+            else:
+                # Ensure existing state tensors are on the correct device
+                if self.state['avg_grad'].device != grad_sample.device:
+                    self.state['avg_grad'] = self.state['avg_grad'].to(device=grad_sample.device)
+                    self.state['avg_nxg'] = self.state['avg_nxg'].to(device=grad_sample.device)
+                    self.state['avg_gsq'] = self.state['avg_gsq'].to(device=grad_sample.device)
+            
+            # Update running averages
             self.state["avg_grad"] = _welford_mean(
                 self.state["avg_grad"], grad_sample, count
             )
@@ -165,22 +219,58 @@ def _restore_param_average(
             elif self.hess_approx == 'gradsq':
                 self.state['avg_gsq'] = _welford_mean(
                     self.state['avg_gsq'], grad_sample.square(), count)
+            
+            # Ensure current_step is incremented
+            self.current_step = count
 
     @torch.no_grad()
     def step(self, closure: ClosureType = None) -> Optional[Tensor]:
+        # Ensure buffers are on correct device
+        self._ensure_buffers_on_device()
+        
+        # Compute total number of parameters
+        total_numel = sum(group['numel'] for group in self.param_groups)
+        
+        # Ensure state is properly initialized
+        if 'count' not in self.state or self.state['count'] is None:
+            self.state['count'] = 0
+        
+        # Initialize gradient-related states if not already done
+        if self.state.get('avg_grad') is None:
+            # Initialize with parameter dtype
+            self.state['avg_grad'] = torch.zeros(
+                total_numel, 
+                device=self._device, 
+                dtype=self._dtype
+            )
+            self.state['avg_nxg'] = torch.zeros_like(self.state['avg_grad'])
+            self.state['avg_gsq'] = torch.zeros_like(self.state['avg_grad'])
+        
+        # Set closure default if not provided
         if closure is None:
             loss = None
         else:
+            # Handle multiple Monte Carlo samples if specified
             losses = []
             for _ in range(self.mc_samples):
                 with torch.enable_grad():
                     loss = closure()
                 losses.append(loss)
             loss = sum(losses) / self.mc_samples
-        if self.sync and dist.is_initialized():  # explicit sync
+        
+        # Handle distributed sync if needed
+        if self.sync and dist.is_initialized():
             self._sync_samples()
+        
+        # Track current step
+        self.current_step += 1
+        
+        # Perform parameter update
         self._update()
-        self._reset_samples()
+        
+        # Reset running statistics, but keep count and state
+        self._reset_samples(reset_count=False)
+        
         return loss
 
     def _sync_samples(self):
@@ -191,17 +281,21 @@ def _sync_samples(self):
         self.state["avg_nxg"].div_(world_size)
 
     def _sample_params(self) -> Tuple[Tensor, Tensor]:
+        # Ensure buffers are on correct device
+        self._ensure_buffers_on_device()
+        
         noise_samples = []
         param_avgs = []
 
         offset = 0
         for group in self.param_groups:
             gnumel = group["numel"]
+            
+            # Generate noise on the same device as the hess buffer
+            hess_buffer = group["hess"] + group["weight_decay"]
             noise_sample = (
-                torch.randn(gnumel, device=self._device, dtype=self._dtype)
-                / (
-                    group["ess"] * (group["hess"] + group["weight_decay"])
-                ).sqrt()
+                torch.randn(gnumel, device=hess_buffer.device, dtype=self._dtype)
+                / (group["ess"] * hess_buffer).sqrt()
             )
             noise_samples.append(noise_sample)
 
@@ -215,7 +309,10 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]:
                 p_noise = noise_sample[goffset : goffset + numel]
 
                 param_avgs.append(p_avg)
-                p.data = (p_avg + p_noise).view(p.shape)
+                # Use autocast for parameter sampling arithmetic
+                with torch.autocast(device_type=p.device.type):
+                    sampled_param = (p_avg + p_noise).to(dtype=p.dtype)
+                p.data = sampled_param.view(p.shape)
                 goffset += numel
                 offset += numel
             assert goffset == group["numel"]  # sanity check
@@ -224,8 +321,20 @@ def _sample_params(self) -> Tuple[Tensor, Tensor]:
         return torch.cat(param_avgs, 0), torch.cat(noise_samples, 0)
 
     def _update(self):
-        self.current_step += 1
+        # If avg_grad is None, initialize with zeros
+        if self.state.get('avg_grad') is None:
+            # Compute total number of parameters
+            total_numel = sum(group['numel'] for group in self.param_groups)
+            self.state['avg_grad'] = torch.zeros(
+                total_numel, 
+                device=self._device, 
+                dtype=self._dtype
+            )
+            self.state['avg_nxg'] = torch.zeros_like(self.state['avg_grad'])
+            self.state['avg_gsq'] = torch.zeros_like(self.state['avg_grad'])
+            self.state['count'] = 0
 
+        # Prepare for parameter update
         offset = 0
         for group in self.param_groups:
             lr = group["lr"]
@@ -237,39 +346,44 @@ def _update(self):
                 [p.flatten() for p in group["params"] if p is not None], 0
             )
 
-            group["momentum"] = self._new_momentum(
-                self.state["avg_grad"][pg_slice], group["momentum"], b1
-            )
-
-            group["hess"] = self._new_hess(
-                self.hess_approx,
-                group["hess"],
-                self.state["avg_nxg"],
-                self.state['avg_gsq'],
-                pg_slice,
-                group["ess"],
-                b2,
-                group["weight_decay"],
-            )
-
-            param_avg = self._new_param_averages(
-                param_avg,
-                group["hess"],
-                group["momentum"],
-                lr * (group["hess_init"] + group["weight_decay"]) if self.rescale_lr else lr,
-                group["weight_decay"],
-                group["clip_radius"],
-                1.0 - pow(b1, float(self.current_step)) if self.debias else 1.0,
-                group["hess_init"]
-            )
+            # Handle case where avg_grad might not have been collected
+            if self.state['avg_grad'] is not None:
+                # Update momentum with available gradient information
+                group["momentum"] = self._new_momentum(
+                    self.state["avg_grad"][pg_slice], group["momentum"], b1
+                )
+
+                group["hess"] = self._new_hess(
+                    self.hess_approx,
+                    group["hess"],
+                    self.state["avg_nxg"],
+                    self.state['avg_gsq'],
+                    pg_slice,
+                    group["ess"],
+                    b2,
+                    group["weight_decay"],
+                )
+
+                param_avg = self._new_param_averages(
+                    param_avg,
+                    group["hess"],
+                    group["momentum"],
+                    lr * (group["hess_init"] + group["weight_decay"]) if self.rescale_lr else lr,
+                    group["weight_decay"],
+                    group["clip_radius"],
+                    1.0 - pow(b1, float(self.current_step)) if self.debias else 1.0,
+                    group["hess_init"]
+                )
 
             # update params
             pg_offset = 0
             for p in group["params"]:
                 if p is not None:
-                    p.data = param_avg[pg_offset : pg_offset + p.numel()].view(
+                    # Convert back to parameter's original dtype
+                    new_param_data = param_avg[pg_offset : pg_offset + p.numel()].view(
                         p.shape
-                    )
+                    ).to(dtype=p.dtype)
+                    p.data = new_param_data
                     pg_offset += p.numel()
             assert pg_offset == group["numel"]  # sanity check
             offset += group["numel"]
diff --git a/requirements-test.txt b/requirements-test.txt
new file mode 100644
index 0000000..9db01bc
--- /dev/null
+++ b/requirements-test.txt
@@ -0,0 +1,3 @@
+pytest >=8.1.1
+numpy < 2.0
+accelerate
diff --git a/tests/test_ivon.py b/tests/test_ivon.py
new file mode 100644
index 0000000..103ef84
--- /dev/null
+++ b/tests/test_ivon.py
@@ -0,0 +1,528 @@
+import torch
+import pytest
+from ivon import IVON
+from accelerate import Accelerator
+
+
+class SimpleNet(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = torch.nn.Linear(10, 2)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+@pytest.fixture
+def model_and_data():
+    torch.manual_seed(42)
+    model = SimpleNet()
+    X = torch.randn(20, 10)
+    y = torch.randint(0, 2, (20,))
+    return model, X, y
+
+
+def test_ivon_initialization():
+    """Test basic IVON optimizer initialization"""
+    model = SimpleNet()
+    optimizer = IVON(
+        model.parameters(),
+        lr=0.1,
+        ess=20,  # effective sample size
+        mc_samples=1,
+    )
+
+    # Check key attributes are set correctly
+    assert optimizer.mc_samples == 1
+    assert optimizer.current_step == 0
+    assert optimizer.hess_approx == "price"
+
+
+def test_ivon_with_accelerate_gradient_accumulation(model_and_data):
+    """Test IVON optimizer with Accelerate gradient accumulation"""
+    model, X, y = model_and_data
+
+    # Move input data to a consistent device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    X = X.to(device)
+    y = y.to(device)
+    model = model.to(device)
+
+    # Accelerate setup
+    accelerator = Accelerator(gradient_accumulation_steps=2)
+
+    criterion = torch.nn.CrossEntropyLoss()
+    base_optimizer = IVON(model.parameters(), lr=0.1, ess=len(X), mc_samples=1)
+
+    # Create dataloader
+    dataloader = torch.utils.data.DataLoader(list(zip(X, y)), batch_size=2)
+
+    # Prepare with Accelerate
+    model, optimizer, dataloader = accelerator.prepare(
+        model, base_optimizer, dataloader
+    )
+
+    model.train()
+
+    # Store initial parameters and state
+    initial_params = [p.clone() for p in model.parameters()]
+
+    base_opt = optimizer.optimizer
+
+    initial_state = {
+        "count": base_opt.current_step,
+        "avg_grad": base_opt.state.get("avg_grad", None).clone()
+        if base_opt.state.get("avg_grad") is not None
+        else None,
+        "avg_nxg": base_opt.state.get("avg_nxg", None).clone()
+        if base_opt.state.get("avg_nxg") is not None
+        else None,
+        "avg_gsq": base_opt.state.get("avg_gsq", None).clone()
+        if base_opt.state.get("avg_gsq") is not None
+        else None,
+    }
+
+    # Keep track of parameters
+    batch_step_count = 0
+
+    for _ in range(2):
+        for _, (batch_x, batch_y) in enumerate(dataloader):
+            with accelerator.accumulate(model):
+                with optimizer.optimizer.sampled_params(train=True):
+                    # Forward pass
+                    outputs = model(batch_x)
+                    loss = criterion(outputs, batch_y)
+
+                    # Backward pass
+                    accelerator.backward(loss)
+
+                # Call step when sync_gradients is True
+                if accelerator.sync_gradients:
+                    optimizer.step()
+                    batch_step_count += 1
+
+    # Check that parameters have been updated
+    updated_params = list(model.parameters())
+    any_param_updated = any(
+        not torch.equal(init_p, updated_p)
+        for init_p, updated_p in zip(initial_params, updated_params)
+    )
+    assert any_param_updated, "Parameters were not updated during training"
+
+    # Check state has been updated
+    current_count = base_opt.current_step
+    initial_count = initial_state["count"]
+
+    # Ensure step has been called
+    assert current_count > initial_count, "Optimizer step count not incremented"
+
+    # Check grad-related states have been updated
+    final_avg_grad = base_opt.state.get("avg_grad")
+    final_avg_nxg = base_opt.state.get("avg_nxg")
+    final_avg_gsq = base_opt.state.get("avg_gsq")
+
+    assert final_avg_grad is not None, "avg_grad not initialized"
+    assert final_avg_nxg is not None, "avg_nxg not initialized"
+    if base_opt.hess_approx == "gradsq":
+        assert final_avg_gsq is not None, "avg_gsq not initialized"
+
+
+def test_pytorch_native_gradient_accumulation(model_and_data):
+    """Test IVON optimizer with PyTorch native gradient accumulation"""
+    model, X, y = model_and_data
+
+    # Move input data to a consistent device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    X = X.to(device)
+    y = y.to(device)
+    model = model.to(device)
+
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = IVON(model.parameters(), lr=0.1, ess=len(X), mc_samples=1)
+
+    # Gradient accumulation parameters
+    accumulation_steps = 2
+
+    # Store initial parameters and state
+    initial_params = [p.clone() for p in model.parameters()]
+    initial_state = {
+        "count": optimizer.state.get("count", None),
+        "avg_grad": optimizer.state.get("avg_grad", None),
+        "avg_nxg": optimizer.state.get("avg_nxg", None),
+        "avg_gsq": optimizer.state.get("avg_gsq", None),
+    }
+
+    model.train()
+    running_loss = 0.0
+
+    for _ in range(2):
+        for i in range(len(X) // 2):
+            # Select a subset of data
+            batch_x = X[i * 2 : (i + 1) * 2]
+            batch_y = y[i * 2 : (i + 1) * 2]
+
+            with optimizer.sampled_params(train=True):
+                # Forward pass
+                outputs = model(batch_x)
+                loss = criterion(outputs, batch_y) / accumulation_steps
+
+                # Backward pass
+                loss.backward()
+                running_loss += loss.item()
+
+            # Update weights only after accumulation_steps
+            if (i + 1) % accumulation_steps == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+                running_loss = 0.0
+
+    # Check that parameters have been updated
+    updated_params = list(model.parameters())
+    any_param_updated = any(
+        not torch.equal(init_p, updated_p)
+        for init_p, updated_p in zip(initial_params, updated_params)
+    )
+    assert any_param_updated, "Parameters were not updated during training"
+
+    # Check state has been updated
+
+    # Instead of checking state['count'], check current_step
+    assert optimizer.current_step > initial_state["count"], (
+        "Optimizer step count not incremented"
+    )
+
+    # Check grad-related states have been updated (they start as None, so should not be None now)
+    base_opt = optimizer
+
+    def validate_state(state_name):
+        # Get the state value
+        state_value = base_opt.state.get(state_name)
+
+        # Validate the state value
+        assert state_value is not None, f"{state_name} not initialized"
+        assert not torch.all(state_value == 0), f"{state_name} is a zero tensor"
+        assert state_value.numel() > 0, f"{state_name} is empty"
+
+    validate_state("avg_grad")
+    validate_state("avg_nxg")
+
+    if base_opt.hess_approx == "gradsq":
+        validate_state("avg_gsq")
+
+
+def test_ivon_sampling(model_and_data):
+    """Test parameter sampling functionality"""
+    model, X, y = model_and_data
+
+    optimizer = IVON(model.parameters(), lr=0.1, ess=len(X), mc_samples=3)
+
+    # Store original parameters
+    original_params = [p.clone() for p in model.parameters()]
+
+    # Collect sampled parameters
+    sampled_params_collection = []
+
+    for _ in range(5):
+        with optimizer.sampled_params():
+            # Collect current parameter values
+            current_params = [p.clone() for p in model.parameters()]
+            sampled_params_collection.append(current_params)
+
+    # Verify that sampling introduces variation
+    variations = [
+        not torch.equal(original_params[i], sampled_params_collection[0][i])
+        for i in range(len(original_params))
+    ]
+    assert any(variations), "Parameter sampling did not introduce variations"
+
+
+def test_ivon_hess_approx_methods():
+    """Test different Hessian approximation methods"""
+    model = SimpleNet()
+
+    # Test 'price' method (default)
+    optimizer_price = IVON(model.parameters(), lr=0.1, ess=20, hess_approx="price")
+    assert optimizer_price.hess_approx == "price"
+
+    # Test 'gradsq' method
+    optimizer_gradsq = IVON(model.parameters(), lr=0.1, ess=20, hess_approx="gradsq")
+    assert optimizer_gradsq.hess_approx == "gradsq"
+
+    # Ensure invalid method raises ValueError
+    with pytest.raises(ValueError):
+        IVON(model.parameters(), lr=0.1, ess=20, hess_approx="invalid_method")
+
+
+def test_ivon_device_move():
+    """Test IVON optimizer when model is moved to different device after initialization"""
+    model = SimpleNet()
+
+    # Initialize optimizer on CPU
+    optimizer = IVON(model.parameters(), lr=0.1, ess=20, mc_samples=1)
+
+    # Move model to CUDA if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        model = model.to(device)
+
+        # This should work without device mismatch errors
+        with optimizer.sampled_params():
+            # Collect current parameter values
+            current_params = [p.clone() for p in model.parameters()]
+
+        # Test training step as well
+        X = torch.randn(10, 10).to(device)
+        y = torch.randint(0, 2, (10,)).to(device)
+        criterion = torch.nn.CrossEntropyLoss()
+
+        with optimizer.sampled_params(train=True):
+            optimizer.zero_grad()
+            outputs = model(X)
+            loss = criterion(outputs, y)
+            loss.backward()
+
+        # This should also work
+        optimizer.step()
+    else:
+        pytest.skip("CUDA not available")
+
+
+def test_ivon_device_state_consistency():
+    """Test that IVON maintains device consistency across state tensors"""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    
+    model = SimpleNet()
+    optimizer = IVON(model.parameters(), lr=0.1, ess=20, mc_samples=1)
+    
+    # Initially everything should be on CPU
+    X_cpu = torch.randn(10, 10)
+    y_cpu = torch.randint(0, 2, (10,))
+    criterion = torch.nn.CrossEntropyLoss()
+    
+    # Train one step on CPU to initialize state
+    with optimizer.sampled_params(train=True):
+        optimizer.zero_grad()
+        outputs = model(X_cpu)
+        loss = criterion(outputs, y_cpu)
+        loss.backward()
+    optimizer.step()
+    
+    # Verify state is initialized on CPU
+    assert optimizer.state['avg_grad'].device.type == 'cpu'
+    assert optimizer.state['avg_nxg'].device.type == 'cpu'
+    
+    # Move model to CUDA
+    model = model.cuda()
+    X_cuda = X_cpu.cuda()
+    y_cuda = y_cpu.cuda()
+    
+    # Train another step on CUDA - this should move state tensors automatically
+    with optimizer.sampled_params(train=True):
+        optimizer.zero_grad()
+        outputs = model(X_cuda)
+        loss = criterion(outputs, y_cuda)
+        loss.backward()
+    optimizer.step()
+    
+    # Verify state tensors are now on CUDA
+    assert optimizer.state['avg_grad'].device.type == 'cuda'
+    assert optimizer.state['avg_nxg'].device.type == 'cuda'
+    
+    # Move back to CPU
+    model = model.cpu()
+    
+    # Train another step on CPU - this should move state tensors back to CPU
+    with optimizer.sampled_params(train=True):
+        optimizer.zero_grad()
+        outputs = model(X_cpu)
+        loss = criterion(outputs, y_cpu)
+        loss.backward()
+    optimizer.step()
+    
+    # Verify state tensors are back on CPU
+    assert optimizer.state['avg_grad'].device.type == 'cpu'
+    assert optimizer.state['avg_nxg'].device.type == 'cpu'
+
+
+def test_ivon_mixed_precision():
+    """Test IVON optimizer with mixed precision (bfloat16)"""
+    model = SimpleNet()
+
+    # Convert model to bfloat16
+    model = model.to(dtype=torch.bfloat16)
+
+    optimizer = IVON(model.parameters(), lr=0.1, ess=20, mc_samples=1)
+
+    X = torch.randn(10, 10, dtype=torch.bfloat16)
+    y = torch.randint(0, 2, (10,))
+    criterion = torch.nn.CrossEntropyLoss()
+
+    # Store original parameter dtypes
+    original_dtypes = {id(p): p.dtype for p in model.parameters()}
+
+    # Test multiple training steps to ensure dtype consistency
+    for step in range(3):
+        with optimizer.sampled_params(train=True):
+            optimizer.zero_grad()
+            outputs = model(X)
+            loss = criterion(outputs, y)
+            loss.backward()
+
+        # Check gradients after sampled_params processing
+        for param in model.parameters():
+            if param.grad is not None:
+                assert param.grad.dtype == param.dtype, (
+                    f"Step {step}: Gradient dtype {param.grad.dtype} doesn't match parameter dtype {param.dtype}"
+                )
+
+        # Verify parameters still have correct dtype
+        for param in model.parameters():
+            assert param.dtype == original_dtypes[id(param)], (
+                f"Step {step}: Parameter dtype changed from {original_dtypes[id(param)]} to {param.dtype}"
+            )
+
+        # Step the optimizer
+        optimizer.step()
+
+        # Final check after step
+        for param in model.parameters():
+            if param.grad is not None:
+                assert param.grad.dtype == param.dtype, (
+                    f"Step {step} after step(): Gradient dtype {param.grad.dtype} doesn't match parameter dtype {param.dtype}"
+                )
+
+
+def test_ivon_autocast_mixed_precision():
+    """Test IVON optimizer with autocast mixed precision"""
+    model = SimpleNet()
+    optimizer = IVON(model.parameters(), lr=0.1, ess=20, mc_samples=1)
+
+    X = torch.randn(10, 10)
+    y = torch.randint(0, 2, (10,))
+    criterion = torch.nn.CrossEntropyLoss()
+
+    # Store original parameter dtypes
+    original_dtypes = {id(p): p.dtype for p in model.parameters()}
+
+    # Test training with autocast (simulates mixed precision training)
+    with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
+        with optimizer.sampled_params(train=True):
+            optimizer.zero_grad()
+            outputs = model(X)
+            loss = criterion(outputs, y)
+            loss.backward()
+
+    # Check that gradients maintain consistency
+    for param in model.parameters():
+        if param.grad is not None:
+            # Gradients might be in different precision due to autocast
+            assert param.grad.device == param.device, (
+                f"Gradient device {param.grad.device} doesn't match parameter device {param.device}"
+            )
+
+    # Verify parameters still have correct dtype
+    for param in model.parameters():
+        assert param.dtype == original_dtypes[id(param)], (
+            f"Parameter dtype changed from {original_dtypes[id(param)]} to {param.dtype}"
+        )
+
+    # Step the optimizer
+    optimizer.step()
+
+    # Final check - parameters should maintain their original dtype
+    for param in model.parameters():
+        assert param.dtype == original_dtypes[id(param)], (
+            f"After step: Parameter dtype changed from {original_dtypes[id(param)]} to {param.dtype}"
+        )
+
+
+def test_ivon_mixed_precision_consistency():
+    """Test IVON optimizer maintains dtype consistency in mixed precision scenarios"""
+    model = SimpleNet()
+
+    # Start with bfloat16 model
+    model = model.to(dtype=torch.bfloat16)
+
+    optimizer = IVON(model.parameters(), lr=0.1, ess=20, mc_samples=1)
+
+    X = torch.randn(10, 10, dtype=torch.bfloat16)
+    y = torch.randint(0, 2, (10,))
+    criterion = torch.nn.CrossEntropyLoss()
+
+    # Multiple training steps to test consistency
+    for step in range(5):
+        with optimizer.sampled_params(train=True):
+            optimizer.zero_grad()
+            outputs = model(X)
+            loss = criterion(outputs, y)
+            loss.backward()
+
+        # Check gradients and parameters maintain consistency
+        for param in model.parameters():
+            assert param.dtype == torch.bfloat16, (
+                f"Step {step}: Parameter dtype changed to {param.dtype}"
+            )
+            if param.grad is not None:
+                # Note: gradients might be promoted for stability, but parameters should maintain their dtype
+                assert param.device == param.grad.device, (
+                    f"Step {step}: Device mismatch"
+                )
+
+        optimizer.step()
+
+        # After step, parameters should still be bfloat16
+        for param in model.parameters():
+            assert param.dtype == torch.bfloat16, (
+                f"Step {step} after step(): Parameter dtype changed to {param.dtype}"
+            )
+
+
+def test_ivon_complex_device_scenario():
+    """Test IVON optimizer with complex device movement scenarios"""
+    model = SimpleNet()
+    
+    # Initialize optimizer on CPU
+    optimizer = IVON(model.parameters(), lr=0.1, ess=20, mc_samples=1)
+    
+    # Move model to CUDA if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        # Move model to CUDA after optimizer initialization
+        model = model.to(device)
+        
+        # Create data on CUDA
+        X = torch.randn(10, 10).to(device)
+        y = torch.randint(0, 2, (10,)).to(device)
+        criterion = torch.nn.CrossEntropyLoss()
+        
+        # Test multiple training steps with device movement
+        for step in range(3):
+            # This should work without device errors
+            with optimizer.sampled_params(train=True):
+                optimizer.zero_grad()
+                outputs = model(X)
+                loss = criterion(outputs, y)
+                loss.backward()
+            
+            optimizer.step()
+            
+            # Verify all parameters are on the same device
+            for param in model.parameters():
+                assert param.device.type == device.type, f"Parameter not on {device.type}"
+        
+        # Test moving model back to CPU
+        model = model.to("cpu")
+        X = X.to("cpu")
+        y = y.to("cpu")
+        
+        # This should also work
+        with optimizer.sampled_params(train=True):
+            optimizer.zero_grad()
+            outputs = model(X)
+            loss = criterion(outputs, y)
+            loss.backward()
+        
+        optimizer.step()
+    else:
+        pytest.skip("CUDA not available")
diff --git a/tests/test_ivon_robustness.py b/tests/test_ivon_robustness.py
new file mode 100644
index 0000000..a879679
--- /dev/null
+++ b/tests/test_ivon_robustness.py
@@ -0,0 +1,322 @@
+import torch
+import numpy as np
+import pytest
+from ivon import IVON
+
+
+class RegressionNet(torch.nn.Module):
+    def __init__(self, input_dim=10, output_dim=1):
+        super().__init__()
+        self.fc1 = torch.nn.Linear(input_dim, 50)
+        self.fc2 = torch.nn.Linear(50, output_dim)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        x = self.relu(self.fc1(x))
+        return self.fc2(x)
+
+
+def generate_regression_data(n_samples=200, input_dim=10, noise_std=0.1):
+    """Generate synthetic regression data"""
+    torch.manual_seed(42)
+    np.random.seed(42)
+
+    # True underlying function (non-linear)
+    true_weights1 = np.random.randn(input_dim, 50)
+    true_weights2 = np.random.randn(50, 1)
+    true_bias1 = np.random.randn(50)
+    true_bias2 = np.random.randn(1)
+
+    X = torch.randn(n_samples, input_dim)
+
+    # Non-linear transformation
+    hidden = torch.relu(
+        torch.matmul(X, torch.tensor(true_weights1, dtype=torch.float32))
+        + torch.tensor(true_bias1, dtype=torch.float32)
+    )
+    y = torch.matmul(
+        hidden, torch.tensor(true_weights2, dtype=torch.float32)
+    ) + torch.tensor(true_bias2, dtype=torch.float32)
+
+    # Add some noise
+    y += torch.randn_like(y) * noise_std
+
+    return X, y
+
+
+def test_ivon_sampling_variance():
+    """Test that IVON introduces parameter sampling variation"""
+    model = RegressionNet()
+    optimizer = IVON(model.parameters(), lr=0.01, ess=100, mc_samples=5)
+
+    # Store initial parameters
+    initial_params = [p.clone() for p in model.parameters()]
+
+    # Collect sampled parameters
+    sampled_params_collection = []
+
+    for _ in range(10):
+        with optimizer.sampled_params():
+            # Collect current parameter values
+            current_params = [p.clone() for p in model.parameters()]
+            sampled_params_collection.append(current_params)
+
+    # Verify sampling introduces variation
+    def param_variation(orig_params, sampled_params):
+        return any(
+            not torch.equal(orig_p, sampled_p)
+            for orig_p, sampled_p in zip(orig_params, sampled_params)
+        )
+
+    variations = [
+        param_variation(initial_params, sampled_params)
+        for sampled_params in sampled_params_collection
+    ]
+
+    assert any(variations), "IVON sampling did not introduce parameter variations"
+
+
+def test_ivon_convergence_comparison():
+    """Compare IVON optimizer convergence with Adam optimizer"""
+    # Prepare data
+    X, y = generate_regression_data(noise_std=0.5)
+
+    # Setup models
+    def train_model(optimizer_class, **optimizer_kwargs):
+        torch.manual_seed(42)  # Ensure reproducibility
+        model = RegressionNet()
+        criterion = torch.nn.MSELoss()
+
+        # Choose optimizer with adaptive hyperparameters
+        if optimizer_class == IVON:
+            optimizer = optimizer_class(
+                model.parameters(),
+                lr=0.001,
+                ess=len(X),
+                mc_samples=5,
+                weight_decay=1e-4,
+                hess_init=10.0,
+                **optimizer_kwargs,
+            )
+        else:
+            optimizer = optimizer_class(
+                model.parameters(), lr=0.001, weight_decay=1e-4, **optimizer_kwargs
+            )
+
+        # Training loop with early stopping
+        best_loss = float("inf")
+        patience = 20
+        patience_counter = 0
+        losses = []
+
+        for epoch in range(300):
+            # Adaptive learning rate decay
+            if epoch > 0 and epoch % 50 == 0:
+                for param_group in optimizer.param_groups:
+                    param_group["lr"] *= 0.5
+
+            def closure():
+                optimizer.zero_grad()
+                outputs = model(X)
+                loss = criterion(outputs, y)
+                loss.backward()
+                return loss
+
+            # For IVON, use sampled params
+            if optimizer_class == IVON:
+                with optimizer.sampled_params(train=True):
+                    loss = closure()
+                optimizer.step()
+            else:
+                loss = closure()
+                optimizer.step()
+
+            current_loss = loss.item()
+            losses.append(current_loss)
+
+            # Early stopping
+            if current_loss < best_loss * 0.99:
+                best_loss = current_loss
+                patience_counter = 0
+            else:
+                patience_counter += 1
+
+            if patience_counter >= patience:
+                break
+
+        return model, losses
+
+    # Train with different optimizers
+    ivon_model, ivon_losses = train_model(IVON)
+    adam_model, adam_losses = train_model(torch.optim.Adam)
+
+    # Evaluate models
+    def evaluate_model(model, X, y):
+        with torch.no_grad():
+            outputs = model(X)
+            mse = torch.nn.functional.mse_loss(outputs, y)
+            return mse.item()
+
+    ivon_mse = evaluate_model(ivon_model, X, y)
+    adam_mse = evaluate_model(adam_model, X, y)
+
+    # Validate learning progression
+    def validate_learning_progression(losses):
+        loss_array = np.array(losses)
+        loss_reduction_rates = np.polyfit(np.arange(len(loss_array)), loss_array, 1)
+        return loss_reduction_rates[0]  # Slope of the linear fit
+
+    # Compute learning progression
+    ivon_learning_slope = validate_learning_progression(ivon_losses)
+    adam_learning_slope = validate_learning_progression(adam_losses)
+
+    # Ensure meaningful learning is happening
+    assert ivon_learning_slope < 0, "IVON failed to show loss reduction"
+    assert adam_learning_slope < 0, "Adam failed to show loss reduction"
+
+    # Performance validation - IVON should perform reasonably compared to Adam
+    max_acceptable_mse = adam_mse * 10  # Allow IVON to be up to 10x worse than Adam
+    assert ivon_mse <= max_acceptable_mse, (
+        f"IVON optimizer failed to converge within acceptable range (MSE: {ivon_mse}, Max Acceptable: {max_acceptable_mse})"
+    )
+
+    # Deviation check
+    performance_deviation = abs(ivon_mse - adam_mse) / (adam_mse + 1e-7)
+    assert performance_deviation < 5.0, (
+        f"IVON performance deviates too much from Adam (Deviation: {performance_deviation:.4f})"
+    )
+
+
+def test_ivon_uncertainty_estimation():
+    """Test IVON's uncertainty estimation capabilities"""
+    X, y = generate_regression_data()
+
+    model = RegressionNet()
+    optimizer = IVON(model.parameters(), lr=0.01, ess=len(X), mc_samples=10)
+
+    # Train the model
+    criterion = torch.nn.MSELoss()
+    for _ in range(100):
+
+        def closure():
+            optimizer.zero_grad()
+            outputs = model(X)
+            loss = criterion(outputs, y)
+            loss.backward()
+            return loss
+
+        with optimizer.sampled_params(train=True):
+            loss = closure()
+        optimizer.step()
+
+    # Collect multiple predictions
+    sample_predictions = []
+    for _ in range(50):
+        with optimizer.sampled_params():
+            with torch.no_grad():
+                pred = model(X)
+                sample_predictions.append(pred)
+
+    # Convert to numpy for analysis
+    sample_predictions = torch.stack(sample_predictions).numpy()
+
+    # Compute prediction variance
+    pred_variance = np.var(sample_predictions, axis=0)
+
+    # Check variance characteristics
+    assert pred_variance.mean() > 0, "IVON failed to capture prediction uncertainty"
+    assert np.all(pred_variance < 50.0), (
+        "Prediction variance too high"
+    )  # More reasonable bound for variational methods
+
+
+def test_ivon_gradient_accumulation_robustness():
+    """Test IVON optimizer's robustness with gradient accumulation"""
+    X, y = generate_regression_data()
+
+    model = RegressionNet()
+    optimizer = IVON(model.parameters(), lr=0.01, ess=len(X), mc_samples=1)
+
+    # Simulate gradient accumulation manually
+    criterion = torch.nn.MSELoss()
+    running_loss = 0.0
+    accumulation_steps = 4
+
+    initial_params = [p.clone() for p in model.parameters()]
+
+    for epoch in range(2):
+        for i in range(0, len(X), accumulation_steps):
+            batch_x = X[i : i + accumulation_steps]
+            batch_y = y[i : i + accumulation_steps]
+
+            # Simulate gradient accumulation
+            batch_loss = criterion(model(batch_x), batch_y) / accumulation_steps
+            batch_loss.backward()
+            running_loss += batch_loss.item()
+
+            # Step only after accumulation
+            if (i + accumulation_steps) % len(X) == 0:
+
+                def closure():
+                    return running_loss
+
+                optimizer.step()
+                optimizer.zero_grad()
+                running_loss = 0.0
+
+    # Check parameters were updated
+    updated_params = list(model.parameters())
+    any_updated = any(
+        not torch.equal(init_p, updated_p)
+        for init_p, updated_p in zip(initial_params, updated_params)
+    )
+    assert any_updated, "Parameters not updated during gradient accumulation"
+
+
+def test_ivon_multi_device_support():
+    """Test IVON optimizer's multi-device support"""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    # Create models on different devices
+    cpu_model = RegressionNet()
+    cuda_model = RegressionNet().cuda()
+
+    X_cpu = torch.randn(100, 10)
+    y_cpu = torch.randn(100, 1)
+
+    X_cuda = X_cpu.cuda()
+    y_cuda = y_cpu.cuda()
+
+    # Test CPU training
+    cpu_optimizer = IVON(cpu_model.parameters(), lr=0.01, ess=len(X_cpu), mc_samples=1)
+
+    # Test CUDA training
+    cuda_optimizer = IVON(
+        cuda_model.parameters(), lr=0.01, ess=len(X_cuda), mc_samples=1
+    )
+
+    # Train both
+    criterion = torch.nn.MSELoss()
+
+    def train_model(model, optimizer, X, y):
+        for _ in range(50):
+
+            def closure():
+                optimizer.zero_grad()
+                outputs = model(X)
+                loss = criterion(outputs, y)
+                loss.backward()
+                return loss
+
+            with optimizer.sampled_params(train=True):
+                loss = closure()
+            optimizer.step()
+        return model
+
+    # Ensure no exceptions are raised
+    try:
+        train_model(cpu_model, cpu_optimizer, X_cpu, y_cpu)
+        train_model(cuda_model, cuda_optimizer, X_cuda, y_cuda)
+    except Exception as e:
+        pytest.fail(f"Multi-device training failed: {e}")