fix dtype

lsy323 · lsy323 · commit 1ab82dfefa49 · 2024-05-25T02:53:22.000Z
diff --git a/jetstream_pt/layers.py b/jetstream_pt/layers.py
@@ -142,7 +142,7 @@ def forward(self, inputs):
             self.weight,
             (((2,), (1)), ((), ())),
             None,
-            torch.int32,
+            jnp.int32.dtype,
         )
       result = result * self.weight_scaler
       if self.quantize_activation:

Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,7 @@ def forward(self, inputs):`
`142`	`142`	`self.weight,`
`143`	`143`	`(((2,), (1)), ((), ())),`
`144`	`144`	`None,`
`145`		`- torch.int32,`
	`145`	`+ jnp.int32.dtype,`
`146`	`146`	`)`
`147`	`147`	`result = result * self.weight_scaler`
`148`	`148`	`if self.quantize_activation:`