update

CHANGELOG.md

@@ -5,9 +5,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 ### Added
-- Added autograd support for `grouped_matmul` and `segment_matmul` CUDA implementations via `cutlass` ([#61](https://github.com/pyg-team/pyg-lib/pull/61)
 - Added support for PyTorch 1.12 ([#57](https://github.com/pyg-team/pyg-lib/pull/57), [#58](https://github.com/pyg-team/pyg-lib/pull/58))
-- Added `grouped_matmul` and `segment_matmul` CUDA implementations via `cutlass` ([#51](https://github.com/pyg-team/pyg-lib/pull/51), [#56](https://github.com/pyg-team/pyg-lib/pull/56))
+- Added `grouped_matmul` and `segment_matmul` CUDA implementations via `cutlass` ([#51](https://github.com/pyg-team/pyg-lib/pull/51), [#56](https://github.com/pyg-team/pyg-lib/pull/56), [#61](https://github.com/pyg-team/pyg-lib/pull/61))
 - Added `pyg::sampler::neighbor_sample` interface ([#54](https://github.com/pyg-team/pyg-lib/pull/54))
 - Added `pyg::sampler::Mapper` utility for mapping global to local node indices ([#45](https://github.com/pyg-team/pyg-lib/pull/45)))
 - Added benchmark script ([#45](https://github.com/pyg-team/pyg-lib/pull/45))

pyg_lib/csrc/ops/matmul.cpp

@@ -1,15 +1,12 @@
 #include "matmul.h"
+
 #include <ATen/core/dispatch/Dispatcher.h>
 #include <torch/library.h>
 #include <torch/script.h>
 
 namespace pyg {
 namespace ops {
 
-using torch::autograd::AutogradContext;
-using torch::autograd::Variable;
-using torch::autograd::variable_list;
-
 namespace {
 
 std::vector<at::Tensor> _grouped_matmul(const std::vector<at::Tensor>& input,
@@ -35,6 +32,10 @@ at::Tensor _segment_matmul(const at::Tensor& input,
   return op.call(input, ptr, other);
 }
 
+using torch::autograd::AutogradContext;
+using torch::autograd::Variable;
+using torch::autograd::variable_list;
+
 // Performs matrix multiplication according to segments.
 class SegmentMatmul : public torch::autograd::Function<SegmentMatmul> {
  public: