Prepare to return Device instead of String

ankane · Nov 3, 2024 · 5752d8f · 5752d8f
1 parent dde07b9
commit 5752d8f
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 5 deletions.
diff --git a/ext/torch/device.cpp b/ext/torch/device.cpp
@@ -23,5 +23,10 @@ void init_device(Rice::Module& m) {
         std::stringstream s;
         s << self.type();
         return s.str();
+      })
+    .define_method(
+      "_str",
+      [](torch::Device& self) {
+        return self.str();
       });
 }
diff --git a/ext/torch/ext.cpp b/ext/torch/ext.cpp
@@ -31,6 +31,7 @@ void Init_ext()
 
   // keep this order
   init_torch(m);
+  init_device(m);
   init_tensor(m, rb_cTensor, rb_cTensorOptions);
   init_nn(m);
   init_fft(m);
@@ -39,7 +40,6 @@ void Init_ext()
 
   init_backends(m);
   init_cuda(m);
-  init_device(m);
   init_generator(m, rb_cGenerator);
   init_ivalue(m, rb_cIValue);
   init_random(m);

diff --git a/ext/torch/tensor.cpp b/ext/torch/tensor.cpp
@@ -212,11 +212,9 @@ void init_tensor(Rice::Module& m, Rice::Class& c, Rice::Class& rb_cTensorOptions
         return s.str();
       })
     .define_method(
-      "device",
+      "_device",
       [](Tensor& self) {
-        std::stringstream s;
-        s << self.device();
-        return s.str();
+        return self.device();
       })
     .define_method(
       "_data_str",

diff --git a/lib/torch/tensor.rb b/lib/torch/tensor.rb
@@ -209,5 +209,10 @@ def coerce(other)
         raise TypeError, "#{self.class} can't be coerced into #{other.class}"
       end
     end
+
+    # TODO return Device instead of String in 0.19.0
+    def device
+      _device._str
+    end
   end
 end