diff --git a/torchcfm/conditional_flow_matching.py b/torchcfm/conditional_flow_matching.py
index ad8a2e1..3544327 100644
--- a/torchcfm/conditional_flow_matching.py
+++ b/torchcfm/conditional_flow_matching.py
@@ -214,7 +214,15 @@ def compute_lambda(self, t):
         [4] Simulation-free Schrodinger bridges via score and flow matching, Preprint, Tong et al.
         """
         sigma_t = self.compute_sigma_t(t)
-        return 2 * sigma_t / (self.sigma**2 + 1e-8)
+        if self.sigma == 0:
+            # Avoid division by zero when sigma=0 (the default).
+            # When sigma=0 the path is deterministic and lambda is not well-defined;
+            # return ones so downstream loss weighting is a no-op.
+            # Use `t` (always a tensor per the docstring) rather than `sigma_t`,
+            # because the base class's `compute_sigma_t` returns the scalar
+            # `self.sigma` when sigma=0, and `torch.ones_like` requires a tensor.
+            return torch.ones_like(t)
+        return 2 * sigma_t / (self.sigma**2)
 
 
 class ExactOptimalTransportConditionalFlowMatcher(ConditionalFlowMatcher):
@@ -226,17 +234,19 @@ class ExactOptimalTransportConditionalFlowMatcher(ConditionalFlowMatcher):
     It overrides the sample_location_and_conditional_flow.
     """
 
-    def __init__(self, sigma: Union[float, int] = 0.0):
+    def __init__(self, sigma: Union[float, int] = 0.0, ot_method="exact"):
         r"""Initialize the ConditionalFlowMatcher class.
 
         It requires the hyper-parameter $\sigma$.
                 Parameters
                 ----------
                 sigma : Union[float, int]
-                ot_sampler: exact OT method to draw couplings (x0, x1) (see Eq.(17) [1]).
+                ot_method: OT method to draw couplings (x0, x1) (see Eq.(17) [1]).
+                    "exact" computes the exact OT plan; other solvers (e.g.
+                    "sinkhorn") can be passed through to ``OTPlanSampler``.
         """
         super().__init__(sigma)
-        self.ot_sampler = OTPlanSampler(method="exact")
+        self.ot_sampler = OTPlanSampler(method=ot_method)
 
     def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False):
         r"""
diff --git a/torchcfm/models/__init__.py b/torchcfm/models/__init__.py
index da4e87f..68d08f4 100644
--- a/torchcfm/models/__init__.py
+++ b/torchcfm/models/__init__.py
@@ -1 +1 @@
-from .models import MLP
+from .models import MLP, GradModel
diff --git a/torchcfm/optimal_transport.py b/torchcfm/optimal_transport.py
index d380cd9..a5802c5 100644
--- a/torchcfm/optimal_transport.py
+++ b/torchcfm/optimal_transport.py
@@ -86,10 +86,11 @@ def get_map(self, x0, x1):
             M = M / M.max()  # should not be normalized when using minibatches
         p = self.ot_fn(a, b, M.detach().cpu().numpy())
         if not np.all(np.isfinite(p)):
-            print("ERROR: p is not finite")
-            print(p)
-            print("Cost mean, max", M.mean(), M.max())
-            print(x0, x1)
+            warnings.warn(
+                "Non-finite values in OT plan p. "
+                f"p={p}. Cost mean={M.mean()}, max={M.max()}. "
+                f"x0={x0}, x1={x1}."
+            )
         if np.abs(p.sum()) < 1e-8:
             if self.warn:
                 warnings.warn("Numerical errors in OT plan, reverting to uniform plan.")
@@ -279,7 +280,8 @@ def wasserstein(
     ret : float
         Wasserstein distance
     """
-    assert power == 1 or power == 2
+    if power not in (1, 2):
+        raise ValueError(f"power must be 1 or 2, got {power}")
     # ot_fn should take (a, b, M) as arguments where a, b are marginals and
     # M is a cost matrix
     if method == "exact" or method is None:
diff --git a/torchcfm/utils.py b/torchcfm/utils.py
index 9149b1e..63fbaaa 100644
--- a/torchcfm/utils.py
+++ b/torchcfm/utils.py
@@ -1,6 +1,5 @@
 import math
 
-import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from torchdyn.datasets import generate_moons
@@ -9,10 +8,12 @@
 
 
 def eight_normal_sample(n, dim, scale=1, var=1):
+    if dim < 2:
+        raise ValueError(f"dim must be >= 2 (the eight normals are 2D), got {dim}")
     m = torch.distributions.multivariate_normal.MultivariateNormal(
         torch.zeros(dim), math.sqrt(var) * torch.eye(dim)
     )
-    centers = [
+    centers_2d = [
         (1, 0),
         (-1, 0),
         (0, 1),
@@ -22,7 +23,8 @@ def eight_normal_sample(n, dim, scale=1, var=1):
         (-1.0 / np.sqrt(2), 1.0 / np.sqrt(2)),
         (-1.0 / np.sqrt(2), -1.0 / np.sqrt(2)),
     ]
-    centers = torch.tensor(centers) * scale
+    centers = torch.zeros(8, dim)
+    centers[:, :2] = torch.tensor(centers_2d) * scale
     noise = m.sample((n,))
     multi = torch.multinomial(torch.ones(8), n, replacement=True)
     data = []
@@ -54,6 +56,8 @@ def forward(self, t, x, *args, **kwargs):
 
 def plot_trajectories(traj):
     """Plot trajectories of some selected samples."""
+    import matplotlib.pyplot as plt
+
     n = 2000
     plt.figure(figsize=(6, 6))
     plt.scatter(traj[0, :n, 0], traj[0, :n, 1], s=10, alpha=0.8, c="black")