diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ff1d33a..35bf103 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -44,17 +44,19 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-sbt-
 
-      # TEMPORARY: dimwit 0.2-SNAPSHOT is not published to any remote yet, so CI builds it from source.
+      # TEMPORARY: dimwit 0.2-SNAPSHOT is not published to any remote yet, so CI builds it from
       - name: Check out dimwit
         uses: actions/checkout@v4
         with:
-          repository: marcelluethi/dimwit
-          ref: version-0.2-SNAPSHOT
+          repository: dimwit-dev/dimwit
+          ref: main
           path: dimwit
 
+      # The version is pinned here to what build.sbt depends on, so a version bump on dimwit's
+      # main cannot silently leave this build resolving a stale snapshot.
       - name: Publish dimwit-core to the local ivy repo
         working-directory: dimwit
-        run: sbt core/publishLocal
+        run: sbt publishLocal
 
       - name: Check formatting
         run: sbt scalafmtCheckAll
diff --git a/core/src/main/scala/deepwit/base/AffineLayer.scala b/core/src/main/scala/deepwit/base/AffineLayer.scala
index 7f6acc0..aaaa561 100644
--- a/core/src/main/scala/deepwit/base/AffineLayer.scala
+++ b/core/src/main/scala/deepwit/base/AffineLayer.scala
@@ -41,7 +41,7 @@ object AffineLayer:
 
     def identity[In: Λ, V: IsFloating](extent: AxisExtent[In], vtype: VType[V] = VType[Float32]): Params[In, Prime[In], V] =
       Params(
-        weight = Tensor2.eye(extent, vtype),
+        weight = Tensor2(extent).eye(vtype),
         bias = Tensor(Shape(Axis[Prime[In]] -> extent.size), vtype).fill(0f)
       )
 
diff --git a/core/src/main/scala/deepwit/base/LinearLayer.scala b/core/src/main/scala/deepwit/base/LinearLayer.scala
index 4d64d13..ed4cf78 100644
--- a/core/src/main/scala/deepwit/base/LinearLayer.scala
+++ b/core/src/main/scala/deepwit/base/LinearLayer.scala
@@ -42,7 +42,7 @@ object LinearLayer:
     def init[In: Λ, Out: Λ, V: IsFloating](inExtent: AxisExtent[In], outExtent: AxisExtent[Out], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Params[In, Out, V] =
       xavierUniform(inExtent, outExtent, key, vtype, gain)
 
-    def identity[In: Λ, V: IsFloating](extent: AxisExtent[In], vtype: VType[V] = VType[Float32]): Params[In, Prime[In], V] = Params(weight = Tensor2.eye(extent, vtype))
+    def identity[In: Λ, V: IsFloating](extent: AxisExtent[In], vtype: VType[V] = VType[Float32]): Params[In, Prime[In], V] = Params(weight = Tensor2(extent).eye(vtype))
 
     def xavierNormal[In: Λ, Out: Λ, V: IsFloating](inExtent: AxisExtent[In], outExtent: AxisExtent[Out], key: Key, vtype: VType[V] = VType[Float32], gain: Float = 1f): Params[In, Out, V] =
       Params(weight = Init.xavierNormal(inExtent, outExtent, key, vtype, gain = gain))
diff --git a/core/src/test/scala/deepwit/embedder/VocabularyEmbedderSuite.scala b/core/src/test/scala/deepwit/embedder/VocabularyEmbedderSuite.scala
index bec346b..c1593cc 100644
--- a/core/src/test/scala/deepwit/embedder/VocabularyEmbedderSuite.scala
+++ b/core/src/test/scala/deepwit/embedder/VocabularyEmbedderSuite.scala
@@ -32,7 +32,7 @@ class VocabularyEmbedderSuite extends AnyFunSpec with Matchers:
       embedder.unembed(embedding) should approxEqual(Tensor(Shape1(vocabExtent)).fromArray(Array(1f, 3f, 5f)), 1e-6f)
 
     it("recovers the token through unembed when the embeddings are orthonormal"):
-      val orthonormal = VocabularyEmbedder(VocabularyEmbedder.Params(Tensor2.eye(vocabExtent)))
+      val orthonormal = VocabularyEmbedder(VocabularyEmbedder.Params(Tensor2(vocabExtent).eye))
       (0 until 3).foreach: token =>
         val logits = orthonormal.unembed(orthonormal(Tensor0(token)))
         logits.argmax(Axis[A]).item shouldBe token