diff --git a/src/embkit/losses/__init__.py b/src/embkit/losses/__init__.py
index 64c4032..44a070a 100644
--- a/src/embkit/losses/__init__.py
+++ b/src/embkit/losses/__init__.py
@@ -1,4 +1,4 @@
-from .base import VAELoss
+from .base import VAELoss, VAELossOutput
 from .vae_loss import (
     # Concrete nn.Module loss classes
     MSEVAELoss,
diff --git a/src/embkit/losses/base.py b/src/embkit/losses/base.py
index d0c8b46..d62113d 100644
--- a/src/embkit/losses/base.py
+++ b/src/embkit/losses/base.py
@@ -3,7 +3,7 @@
 """
 
 from abc import abstractmethod
-from typing import Tuple
+from typing import Tuple, NamedTuple
 
 import torch
 from torch import nn, Tensor
@@ -64,3 +64,11 @@ def _kl_divergence(mu: Tensor, logvar: Tensor) -> Tensor:
         Returns a 1-D tensor of shape ``(batch_size,)``.
         """
         return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
+
+
+class VAELossOutput(NamedTuple):
+    """Named output of :meth:`VAELoss.forward`."""
+    total: Tensor
+    recon: Tensor
+    kl: Tensor
+
diff --git a/src/embkit/losses/vae_loss.py b/src/embkit/losses/vae_loss.py
index 0f3c42d..63e37d7 100644
--- a/src/embkit/losses/vae_loss.py
+++ b/src/embkit/losses/vae_loss.py
@@ -20,15 +20,9 @@
 import torch.nn.functional as F
 from torch import Tensor
 
-from .base import VAELoss
+from .base import VAELoss, VAELossOutput
 
 
-class VAELossOutput(NamedTuple):
-    """Named output of :meth:`VAELoss.forward`."""
-    total: Tensor
-    recon: Tensor
-    kl: Tensor
-
 
 # ---------------------------------------------------------------------------
 # Concrete loss classes
diff --git a/src/embkit/models/vae/__init__.py b/src/embkit/models/vae/__init__.py
index 628c7f4..48be1f5 100644
--- a/src/embkit/models/vae/__init__.py
+++ b/src/embkit/models/vae/__init__.py
@@ -2,8 +2,8 @@
 This file initializes the Variational Autoencoder (VAE) models in the embkit library.
 It includes two classes, BaseVAE and NetVae, which are used for different types of VAEs respectively.
 """
-from .vae import VAE, BaseVAE
+from .vae import VAE, BaseVAE, VAEOutput
 from .net_vae import NetVAE
 from .rna_vae import RNAVAE
-from .encoder import  Encoder
+from .encoder import  Encoder, VAEEncoder, EncoderOutput
 from .decoder import  Decoder
diff --git a/src/embkit/models/vae/encoder.py b/src/embkit/models/vae/encoder.py
index e5eedca..33884b6 100644
--- a/src/embkit/models/vae/encoder.py
+++ b/src/embkit/models/vae/encoder.py
@@ -40,7 +40,7 @@ class VAEEncoder(nn.Module):
     VAE Encoder wrapper that takes a backbone network (e.g., FFN, CNN), 
     projects features to latent mean and log-variance, and calculates KL divergence.
     """
-    def __init__(self, backbone: nn.Module, feature_dim: int, latent_dim: int):
+    def __init__(self, backbone: nn.Module, feature_dim: int, latent_dim: int, device=None, dtype=None):
         """
         Args:
             backbone (nn.Module): Feature extractor module outputting a tensor of shape (batch_size, feature_dim).
@@ -49,10 +49,12 @@ def __init__(self, backbone: nn.Module, feature_dim: int, latent_dim: int):
         """
         super().__init__()
         self.backbone = backbone
-        
+        self.feature_dim = feature_dim
+        self.latent_dim = latent_dim
+
         # Linear projections for mu and log-variance
-        self.fc_mu = nn.Linear(feature_dim, latent_dim)
-        self.fc_logvar = nn.Linear(feature_dim, latent_dim)
+        self.fc_mu = nn.Linear(feature_dim, latent_dim, device=device, dtype=dtype)
+        self.fc_logvar = nn.Linear(feature_dim, latent_dim, device=device, dtype=dtype)
 
     def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
         """
diff --git a/src/embkit/optimize/__init__.py b/src/embkit/optimize/__init__.py
index fdbbc10..b599c79 100644
--- a/src/embkit/optimize/__init__.py
+++ b/src/embkit/optimize/__init__.py
@@ -328,10 +328,7 @@ def vae_step(batch, beta_value: float) -> Dict[str, torch.Tensor]:
         # nn.Module-based loss: update beta state then call forward
         if beta_value is not None:
             criterion.beta = beta_value
-<<<<<<< HEAD
-=======
-        # print(res, x_tensor, res.mu, res.logvar)
->>>>>>> 3ea8888 (Fixing unit test that broken when positional encoding was updated)
+
         total_loss, recon_loss, kl_loss = criterion(res.recon, x_tensor, res.mu, res.logvar)
 
         return {