diff --git a/stable_pretraining/methods/visreg.py b/stable_pretraining/methods/visreg.py
index f8cfe2556..caa235b15 100644
--- a/stable_pretraining/methods/visreg.py
+++ b/stable_pretraining/methods/visreg.py
@@ -85,13 +85,13 @@ def __init__(
         self._cached_B = -1
         self._cached_target = None
 
-    def _get_target(self, B: int, device, dtype) -> torch.Tensor:
+    def _get_target(self, B: int, device) -> torch.Tensor:
         """Theoretical standard-normal quantiles for ``B`` sorted samples."""
         if self._cached_B != B:
             q = torch.linspace(1, B, B, device=device, dtype=torch.float32) / (B + 1)
             self._cached_target = torch.erfinv(2 * q - 1).mul_(math.sqrt(2))
             self._cached_B = B
-        return self._cached_target.to(device=device, dtype=dtype)
+        return self._cached_target.to(device=device)
 
     def forward(self, z: torch.Tensor) -> torch.Tensor:
         """:param z: Embeddings [V, B, D] (views, batch, dim).
@@ -104,13 +104,13 @@ def forward(self, z: torch.Tensor) -> torch.Tensor:
         center_loss = mu.pow(2).mean()
 
         z_centered = z - mu
-        std = z_centered.norm(dim=1).div(math.sqrt(B)) + 1e-6
+        std = z_centered.norm(dim=1).div(math.sqrt(B)).clamp(min=1e-6)
         scale_loss = (std - 1.0).pow(2).mean()
 
         z_norm = z_centered / std.detach().unsqueeze(1)
         W = F.normalize(torch.randn(D, self.K, device=z.device, dtype=z.dtype), dim=0)
         p_sorted = (z_norm @ W).sort(dim=1).values
-        target = self._get_target(B, z.device, z.dtype).view(1, B, 1)
+        target = self._get_target(B, z.device).view(1, B, 1)
         shape_loss = (p_sorted - target).pow(2).mean()
 
         return (
@@ -207,7 +207,7 @@ def __init__(
         lambda_scale: float = 1.0,
         lambda_shape: float = 1.0,
         lambda_center: float = 1.0,
-        lamb: float = 0.02,
+        lamb: float = 0.9,
         pretrained: bool = False,
         drop_path_rate: float = 0.1,
     ):