Fixed index method for Device when no index [skip ci]

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

Hide file tree

Showing 4 changed files with 16 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 - Improved `inspect` for `Device`
 - Fixed equality for `Device`
+- Fixed `index` method for `Device` when no index
 
 ## 0.18.0 (2024-10-22)
 

diff --git a/ext/torch/device.cpp b/ext/torch/device.cpp
@@ -8,7 +8,7 @@ void init_device(Rice::Module& m) {
   Rice::define_class_under<torch::Device>(m, "Device")
     .define_constructor(Rice::Constructor<torch::Device, const std::string&>())
     .define_method(
-      "index",
+      "_index",
       [](torch::Device& self) {
         return self.index();
       })

diff --git a/lib/torch/device.rb b/lib/torch/device.rb
@@ -1,7 +1,11 @@
 module Torch
   class Device
+    def index
+      index? ? _index : nil
+    end
+
     def inspect
-      extra = ", index: #{index.inspect}" if index != -1
+      extra = ", index: #{index.inspect}" if index?
       "device(type: #{type.inspect}#{extra})"
     end
     alias_method :to_s, :inspect

diff --git a/test/device_test.rb b/test/device_test.rb
@@ -1,6 +1,15 @@
 require_relative "test_helper"
 
 class DeviceTest < Minitest::Test
+  def test_type
+    assert_equal Torch.device("cpu:0").type, "cpu"
+  end
+
+  def test_index
+    assert_nil Torch.device("cpu").index
+    assert_equal 0, Torch.device("cpu:0").index
+  end
+
   def test_equal
     a = Torch.device("cpu")
     b = Torch.device("cpu")