Update to PyTorch 2.1.0

sbrunk · Oct 29, 2023 · 72d7e30 · 72d7e30
1 parent 05078a8
commit 72d7e30
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 8 deletions.
diff --git a/build.sbt b/build.sbt
@@ -24,8 +24,8 @@ ThisBuild / tlSitePublishBranch := Some("main")
 ThisBuild / apiURL := Some(new URL("https://storch.dev/api/"))
 
 val scrImageVersion = "4.0.34"
-val pytorchVersion = "2.0.1"
-val cudaVersion = "12.1-8.9"
+val pytorchVersion = "2.1.0"
+val cudaVersion = "12.3-8.9"
 val openblasVersion = "0.3.23"
 val mklVersion = "2023.1"
 ThisBuild / scalaVersion := "3.3.1"

diff --git a/core/src/main/scala/torch/internal/NativeConverters.scala b/core/src/main/scala/torch/internal/NativeConverters.scala
@@ -24,6 +24,7 @@ import org.bytedeco.pytorch.{
   DeviceOptional,
   DoubleOptional,
   BoolOptional,
+  LongArrayRefOptional,
   LongOptional,
   TensorOptional
 }
@@ -60,6 +61,8 @@ private[torch] object NativeConverters:
           pytorch.ScalarOptional(scalar)
       )
 
+  extension (i: Int) def toScalarOptional = pytorch.ScalarOptional(pytorch.Scalar(i))
+
   extension [D <: DType](t: Tensor[D] | Option[Tensor[D]])
     def toOptional: TensorOptional =
       convertToOptional(t, t => pytorch.TensorOptional(t.native))

diff --git a/core/src/main/scala/torch/ops/PointwiseOps.scala b/core/src/main/scala/torch/ops/PointwiseOps.scala
@@ -631,7 +631,7 @@ private[torch] trait PointwiseOps {
     *
     * @group pointwise_ops
     */
-  def nanToNum[D <: RealNN](
+  def nanToNum[D <: DType](
       input: Tensor[D],
       nan: Option[Double] = None,
       posinf: Option[Double] = None,

diff --git a/core/src/main/scala/torch/ops/ReductionOps.scala b/core/src/main/scala/torch/ops/ReductionOps.scala
@@ -683,7 +683,7 @@ private[torch] trait ReductionOps {
       torchNative.std(
         input.native,
         dim.toArray,
-        correction.toOptional,
+        correction.toScalarOptional,
         keepdim
       )
     )
@@ -722,7 +722,7 @@ private[torch] trait ReductionOps {
       torchNative.std_mean(
         input.native,
         dim.toArray,
-        correction.toOptional,
+        correction.toScalarOptional,
         keepdim
       )
     (fromNative[D](nativeTuple.get0), fromNative[D](nativeTuple.get1))
@@ -828,7 +828,7 @@ private[torch] trait ReductionOps {
       torchNative.`var`(
         input.native,
         dim.toArray,
-        correction.toOptional,
+        correction.toScalarOptional,
         keepdim
       )
     )
@@ -867,7 +867,7 @@ private[torch] trait ReductionOps {
       torchNative.var_mean(
         input.native,
         dim.toArray,
-        correction.toOptional,
+        correction.toScalarOptional,
         keepdim
       )
     (fromNative[D](nativeTuple.get0), fromNative[D](nativeTuple.get1))

diff --git a/core/src/main/scala/torch/special/package.scala b/core/src/main/scala/torch/special/package.scala
@@ -44,7 +44,7 @@ package object special:
     fromNative(torchNative.exp2(input.native))
 
   /** Computes the exponential of the elements minus 1 of `input`. */
-  def expm1[D <: RealNN](input: Tensor[D]): Tensor[FloatPromoted[D]] =
+  def expm1[D <: DType](input: Tensor[D]): Tensor[FloatPromoted[D]] =
     fromNative(torchNative.expm1(input.native))
 
   /** Computes the zeroth order modified Bessel function of the first kind for each element of