torch-linear-assignment 0.0.1.post3__tar.gz → 0.0.2__tar.gz

{torch_linear_assignment-0.0.1.post3/torch_linear_assignment.egg-info → torch_linear_assignment-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torch-linear-assignment
-Version: 0.0.1.post3
+Version: 0.0.2
 Summary: Batched linear assignment with PyTorch and CUDA.
 Author: Ivan Karpukhin
 Author-email: karpuhini@yandex.ru

{torch_linear_assignment-0.0.1.post3 → torch_linear_assignment-0.0.2}/setup.py

@@ -48,7 +48,7 @@ with open("README.md") as fp:
 if __name__ == '__main__':
     setuptools.setup(
         name="torch-linear-assignment",
-        version="0.0.1.post3",
+        version="0.0.2",
         author="Ivan Karpukhin",
         author_email="karpuhini@yandex.ru",
         description="Batched linear assignment with PyTorch and CUDA.",

{torch_linear_assignment-0.0.1.post3 → torch_linear_assignment-0.0.2}/src/torch_linear_assignment_cuda_kernel.cu

@@ -21,10 +21,10 @@
 typedef unsigned char uint8_t;
 
 
-int SMPCores()
+int SMPCores(int device_index)
 {
   cudaDeviceProp devProp;
-  cudaGetDeviceProperties(&devProp, 0);
+  cudaGetDeviceProperties(&devProp, device_index);
   switch (devProp.major){
   case 2: // Fermi
     if (devProp.minor == 1)
@@ -204,8 +204,11 @@ void solve_cuda_kernel_batch(int bs, int nr, int nc,
 
 
 template <typename scalar_t>
-void solve_cuda_batch(int bs, int nr, int nc,
+void solve_cuda_batch(int device_index,
+                      int bs, int nr, int nc,
                       scalar_t *cost, int *col4row, int *row4col) {
+  cudaSetDevice(device_index);
+
   TORCH_CHECK(std::numeric_limits<scalar_t>::has_infinity, "Data type doesn't have infinity.");
   auto infinity = std::numeric_limits<scalar_t>::infinity();
 
@@ -221,9 +224,9 @@ void solve_cuda_batch(int bs, int nr, int nc,
   thrust::fill(v.begin(), v.end(), (scalar_t) 0);
   thrust::fill(path.begin(), path.end(), -1);
 
-  int blockSize = SMPCores();
+  int blockSize = SMPCores(device_index);
   int gridSize = (bs + blockSize - 1) / blockSize;
-  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
+  at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream(device_index);
   solve_cuda_kernel_batch<<<gridSize, blockSize, 0, stream.stream()>>>(
     bs, nr, nc,
     cost,
@@ -260,12 +263,13 @@ std::vector<torch::Tensor> batch_linear_assignment_cuda(torch::Tensor cost) {
     return {col4row, row4col};
   }
 
-  AT_DISPATCH_FLOATING_TYPES(cost.type(), "solve_cuda_batch", ([&] {
+  AT_DISPATCH_FLOATING_TYPES(cost.scalar_type(), "solve_cuda_batch", [&] {
     solve_cuda_batch<scalar_t>(
+        device.index(),
         sizes[0], sizes[1], sizes[2],
         cost.data<scalar_t>(),
         col4row.data<int>(),
         row4col.data<int>());
-  }));
+  });
   return {col4row, row4col};
 }

{torch_linear_assignment-0.0.1.post3 → torch_linear_assignment-0.0.2/torch_linear_assignment.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torch-linear-assignment
-Version: 0.0.1.post3
+Version: 0.0.2
 Summary: Batched linear assignment with PyTorch and CUDA.
 Author: Ivan Karpukhin
 Author-email: karpuhini@yandex.ru