Repository: pytorch/pytorch

Issue:
Dynamic GRU weight
## 🐛 Bug

For weight norm or pruning, this library needs to support dynamically updating the weight tensors. My might attempt to do so that brought up several interesting stack traces:

## To Reproduce

Run this script:
```python3
import torch

# Constants
size = 16
batch_size = 4
seq_len = 8
device = torch.device('cuda')
input_ = torch.randn(seq_len, batch_size, size).to(device)
hidden = torch.randn(1, batch_size, size).to(device)

gru = torch.nn.GRU(size, size).to(device)

# Update weight with a `torch.tensor`
# NOTE: Similar weight update as torch.nn.utils.weight_nrom
data = gru.weight_hh_l0.data
del gru._parameters['weight_hh_l0']
setattr(gru, 'weight_hh_l0', torch.tensor(data))

# Optional call to resolve parameter shapes
gru.flatten_parameters()

# Run forward pass
_, output = gru(input_, hidden)
```

With out ``gru.flatten_parameters``:
```
UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  setattr(gru, 'weight_hh_l0', torch.tensor(data))
Traceback (most recent call last):
  File "ddd.py", line 15, in <module>
    _, output = gru(input_, hidden)
  File "/home/michaelp/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/michaelp/.local/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 179, in forward
    self.dropout, self.training, self.bidirectional, self.batch_first)
RuntimeError: num_ptrs == (num_parameters * (has_biases ? 1 : 2)) ASSERT FAILED at /pytorch/aten/src/ATen/native/cudnn/RNN.cpp:1190, please report a bug to PyTorch.
```

With ``gru.flatten_parameters``:
```
UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  setattr(gru, 'weight_hh_l0', torch.tensor(data))
Traceback (most recent call last):
  File "ddd.py", line 14, in <module>
    gru.flatten_parameters()
  File "/home/michaelp/.local/lib/python3.6/site-packages/torch/nn/modules/rnn.py", line 113, in flatten_parameters
    self.batch_first, bool(self.bidirectional))
RuntimeError: MatrixRef: ArrayRef size 3 not divisible by stride 4
```

## Expected behavior

That I can update the GRU weight with a new `torch.tensor`, without an issue. 

## Environment

```
Collecting environment information...
PyTorch version: 1.0.0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 18.04.1 LTS
GCC version: (Ubuntu 7.3.0-16ubuntu3) 7.3.0
CMake version: version 3.10.2

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla P100-PCIE-16GB
GPU 1: Tesla P100-PCIE-16GB
GPU 2: Tesla P100-PCIE-16GB
GPU 3: Tesla P100-PCIE-16GB

Nvidia driver version: 390.30
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.1.3

Versions of relevant libraries:
[pip] Could not collect
[conda] Could not collect
```

Segfault in dataparallel + checkpoint
Reproducible on master. Reported at https://discuss.pytorch.org/t/segmentation-fault-when-using-checkpoint-and-dataparallel/25247

```py
import torch
import torch.nn as nn
from torch import optim
import torch.utils.checkpoint as chk
import torch.nn.functional as F

class model(nn.Module):

    def __init__(self):
        super(model, self).__init__()

        self.blocks = nn.ModuleDict()

        self.conv0 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1, bias=True)
        self.blocks['0'] = self.conv0

        self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1, bias=True)
        self.blocks['1'] = self.conv1

        self.conv2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1, bias=True)
        self.blocks['2'] = self.conv2

        self.conv3 = nn.Conv2d(64, 20, kernel_size=3, stride=1, padding=1, bias=True)
        self.blocks['3'] = self.conv3

    def forward(self, x):

        x = self.blocks['0'](x)

        #x1 = self.blocks['1'](x)
        x1 = chk.checkpoint(self.conv1,x)

        #x2 = self.blocks['2'](x)
        x2 = chk.checkpoint(self.conv2,x)

        x = torch.cat((x1,x2),1)

        x = self.blocks['3'](x)

        return x
    
test_model = model()
test_model = nn.DataParallel(test_model)
test_model = test_model.cuda()

loss = nn.MSELoss()
optimizer = optim.SGD(test_model.module.parameters(), lr = 0.01)

for i in range(100):
    print(i)
    data = torch.rand(4, 3, 15,15)
    labels = torch.rand(4,20, 15,15).cuda()
    test_preds = test_model(data)
    optimizer.zero_grad()
    test_loss = loss(test_preds, labels)
    test_loss.backward()
    optimizer.step()

print('Finished')
```

Author of the post also provided GDB trace
```
Program received signal SIGSEGV, Segmentation fault.
[Switching to Thread 0x7fff649ff700 (LWP 39737)]
std::__push_heap<__gnu_cxx::__normal_iterator<torch::autograd::FunctionTask*, std::vector<torch::autograd::FunctionTask> >, long, torch::autograd::FunctionTask, __gnu_cxx::__ops::_Iter_comp_val<torch::autograd::CompareFunctionTaskTime> > (__first=..., __holeIndex=1, __topIndex=__topIndex@entry=0, __value=..., __comp=...) at /opt/rh/devtoolset-3/root/usr/include/c++/4.9.2/bits/stl_heap.h:129
129     /opt/rh/devtoolset-3/root/usr/include/c++/4.9.2/bits/stl_heap.h: No such file or directory.
```
torch.eq_() get wrong result in 1.0.0
## 🐛 Bug

<!-- tensor.eq_() behave like "less than", but not "equal"-->

Steps to reproduce the behavior:
>>> a = torch.Tensor([1,2,3])
>>> b = torch.Tensor([2,2,2])
>>> c = a.eq(b) # get correct result in v1.0.0
>>> c
tensor([0, 1, 0], dtype=torch.uint8) # expected behavior
>>> a.eq_(b) # get wrong result in v1.0.0. correct in v0.4.0
tensor([0., 1., 1.])


 - PyTorch Version (e.g., 1.0): stable 1.0 installed from "conda install pytorch torchvision -c pytorch"
 - OS (e.g., Linux): Ubuntu 16.04
 - Python version: python 3.7


Discussion:

Hi, do you guys have any workaround to solve this problem?
Have the same seg fault, both on 0.41 and nightly built. Keep waiting...
+1
> Hi, do you guys have any workaround to solve this problem?

According to the link above, you can warp the whold model to prevent it. Didn't test, I split the model to different gpus though it would be slower...
i think the problem is here
https://github.com/pytorch/pytorch/blob/f52f68bcf9fa490578539fd42c3e050bd5a5f68a/aten/src/ATen/native/LegacyDefinitions.cpp#L114

Resolving patch:
diff --git a/torch/hub.py b/torch/hub.py
--- a/torch/hub.py
+++ b/torch/hub.py
@@ -124,7 +124,6 @@ def load(github, model, force_reload=False, *args, **kwargs):
     # Download zipped code from github
     url = _git_archive_link(repo_info, branch)
     cached_file = os.path.join(hub_dir, branch + '.zip')
-    extracted_repo = os.path.join(hub_dir, repo_name + '-' + branch)
     repo_dir = os.path.join(hub_dir, repo_name + '_' + branch)
 
     use_cache = (not force_reload) and os.path.exists(repo_dir)
@@ -136,13 +135,18 @@ def load(github, model, force_reload=False, *args, **kwargs):
         sys.stderr.write('Using cache found in {}'.format(repo_dir))
     else:
         _remove_if_exists(cached_file)
-        _remove_if_exists(extracted_repo)
-        _remove_if_exists(repo_dir)
-
         _download_url_to_file(url, cached_file)
-        zipfile.ZipFile(cached_file).extractall(hub_dir)
+
+        cached_zipfile = zipfile.ZipFile(cached_file)
+
+        # Github renames folder repo-v1.x.x to repo-1.x.x
+        extraced_repo_name = cached_zipfile.infolist()[0].filename
+        extracted_repo = os.path.join(hub_dir, extraced_repo_name)
+        _remove_if_exists(extracted_repo)
+        cached_zipfile.extractall(hub_dir)
 
         _remove_if_exists(cached_file)
+        _remove_if_exists(repo_dir)
         shutil.move(extracted_repo, repo_dir)  # rename the repo
 
     sys.path.insert(0, repo_dir)  # Make Python interpreter aware of the repo
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -222,7 +222,7 @@ def __setstate__(self, d):
 
     @property
     def _flat_weights(self):
-        return list(self._parameters.values())
+        return [p for layerparams in self.all_weights for p in layerparams]
 
     @property
     def all_weights(self):


Test patch:
diff --git a/aten/src/ATen/test/test_install/CMakeLists.txt b/aten/src/ATen/test/test_install/CMakeLists.txt
--- a/aten/src/ATen/test/test_install/CMakeLists.txt
+++ b/aten/src/ATen/test/test_install/CMakeLists.txt
@@ -3,6 +3,8 @@ find_package(ATen REQUIRED)
 include_directories(${ATEN_INCLUDE_DIR})
 
 # C++11
-set(CMAKE_CXX_FLAGS "--std=c++11 ${CMAKE_CXX_FLAGS}")
+if (not MSVC) 
+    set(CMAKE_CXX_FLAGS "--std=c++11 ${CMAKE_CXX_FLAGS}") 
+endif()
 add_executable(main main.cpp)
 target_link_libraries(main ${ATEN_LIBRARIES})
diff --git a/c10/test/util/intrusive_ptr_test.cpp b/c10/test/util/intrusive_ptr_test.cpp
--- a/c10/test/util/intrusive_ptr_test.cpp
+++ b/c10/test/util/intrusive_ptr_test.cpp
@@ -11,9 +11,11 @@ using c10::intrusive_ptr_target;
 using c10::make_intrusive;
 using c10::weak_intrusive_ptr;
 
+#ifndef _MSC_VER
 #pragma GCC diagnostic ignored "-Wpragmas"
 #pragma GCC diagnostic ignored "-Wunknown-warning-option"
 #pragma GCC diagnostic ignored "-Wself-move"
+#endif
 
 namespace {
 class SomeClass0Parameters : public intrusive_ptr_target {};
diff --git a/test/cpp/api/serialize.cpp b/test/cpp/api/serialize.cpp
--- a/test/cpp/api/serialize.cpp
+++ b/test/cpp/api/serialize.cpp
@@ -215,7 +215,9 @@ TEST(SerializeTest, Optim) {
 TEST(SerializeTest, XOR_CUDA) {
   torch::manual_seed(0);
   // We better be able to save and load a XOR model!
-  auto getLoss = [](Sequential model, uint32_t batch_size, bool is_cuda=false) {
+  auto getLoss = [](Sequential model,
+                    uint32_t batch_size,
+                    bool is_cuda = false) {
     auto inputs = torch::empty({batch_size, 2});
     auto labels = torch::empty({batch_size});
     if (is_cuda) {
@@ -269,3 +271,34 @@ TEST(SerializeTest, XOR_CUDA) {
   loss = getLoss(model3, 100, true);
   ASSERT_LT(loss.item<float>(), 0.1);
 }
+
+TEST(
+    SerializeTest,
+    CanSerializeModulesWithIntermediateModulesWithoutParametersOrBuffers) {
+  struct C : torch::nn::Module {
+    C() {
+      register_buffer("foo", torch::ones(5, torch::kInt32));
+    }
+  };
+  struct B : torch::nn::Module {};
+  struct A : torch::nn::Module {
+    A() {
+      register_module("b", std::make_shared<B>());
+      register_module("c", std::make_shared<C>());
+    }
+  };
+  struct M : torch::nn::Module {
+    M() {
+      register_module("a", std::make_shared<A>());
+    }
+  };
+
+  auto out = std::make_shared<M>();
+  std::stringstream ss;
+  torch::save(out, ss);
+  auto in = std::make_shared<M>();
+  torch::load(in, ss);
+
+  const int output = in->named_buffers()["a.c.foo"].sum().item<int>();
+  ASSERT_EQ(output, 5);
+}
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -14,12 +14,14 @@
 from torch.autograd.gradcheck import gradgradcheck, gradcheck
 from torch.autograd.function import once_differentiable
 from torch.autograd.profiler import profile
+from torch.utils.checkpoint import checkpoint
 from common_utils import (TEST_MKL, TestCase, run_tests, skipIfNoLapack,
                           suppress_warnings, skipIfRocm,
                           prod_single_zero, random_square_matrix_of_rank,
                           random_symmetric_matrix, random_symmetric_psd_matrix,
                           random_symmetric_pd_matrix, make_nonzero_det,
                           random_fullrank_matrix_distinct_singular_value, load_tests)
+from common_cuda import TEST_CUDA
 from torch.autograd import Variable, Function, detect_anomaly
 from torch.autograd.function import InplaceFunction
 from torch.testing import make_non_contiguous, randn_like
@@ -202,6 +204,16 @@ def compute_grad(create_graph):
         x_grad, x_grad_clone = compute_grad(create_graph=True)
         self.assertEqual(x_grad, x_grad_clone)
 
+    def test_sum_to_with_empty_dim_grad(self):
+        a = torch.rand(4, 0, requires_grad=True)
+        b = torch.rand(4, 1, requires_grad=True)
+        c = a + b
+        assert c.shape == (4, 0)
+        c.sum().backward()
+
+        self.assertEqual(b.grad, torch.zeros(4, 1))
+        self.assertEqual(a.grad, torch.zeros(4, 0))
+
     def test_hessian_vector(self):
         x = torch.randn(2, 2, requires_grad=True)
         y = torch.randn(2, 2, requires_grad=True)
@@ -1464,19 +1476,24 @@ def test_ctc_loss(self):
         gradcheck_input_size = 10
 
         # device, input_length
-        tests = [('cpu', 150)]
+        tests = [('cpu', 150, False),
+                 ('cpu', 150, True)]
         if torch.cuda.is_available():
-            tests += [('cuda', 50),
-                      ('cuda', 150)]
+            tests += [('cuda', 50, False),
+                      ('cuda', 150, False),
+                      ('cuda', 50, True),
+                      ('cuda', 150, True)]
 
-        for device, input_length in tests:
+        for device, input_length, vary_lengths in tests:
             targets = torch.randint(1, num_labels, (batch_size, target_length),
                                     device=device, dtype=torch.long)
             x = torch.randn(gradcheck_input_size, device=device, requires_grad=True)
             tile_factors = torch.randn(input_length * batch_size * num_labels // gradcheck_input_size + 1,
                                        device=device)
-            input_lengths = [input_length for _ in range(batch_size)]
-            target_lengths = [target_length for _ in range(batch_size)]
+            input_lengths = [(torch.randint(input_length // 2, input_length + 1, ()).item()
+                              if vary_lengths or i == 0 else input_length) for i in range(batch_size)]
+            target_lengths = [(torch.randint(target_length // 2, target_length + 1, ()).item()
+                               if vary_lengths or i == 0 else target_length) for i in range(batch_size)]
 
             def ctc_after_softmax(x):
                 x_full = ((x[:, None] * tile_factors[None, :]).view(-1)[:input_length * batch_size * num_labels]
@@ -2704,6 +2721,36 @@ def f(inp):
         gradcheck(f, torch.rand(10, dtype=torch.float64, requires_grad=True))
         gradgradcheck(f, torch.rand(10, dtype=torch.float64, requires_grad=True))
 
+    @unittest.skipIf(not TEST_CUDA, "Requires cuda for multi device")
+    def test_multi_device_reentrant_autograd(self):
+        # Output on gpu so that this task will be associated with the gpu thread
+        def fn_on_gpu(inp):
+            # Artificially increase the priority of the next op to make sure it runs
+            # as soon as we reach it before the ops of branch1.
+            dummy = inp * 2 * 2 * 2 * 2
+            return inp.cuda()
+
+        def parent_on_cpu(inp):
+            # Slow branch of ops on gpu so that the work queue for the gpu thread
+            # won't empty too quickly. They also have smaller priorities than the
+            # ones created by fn_on_gpu
+            branch1 = inp.cuda()
+            branch1 = branch1 / branch1
+            branch1 = branch1 / branch1
+            branch1 = branch1 / branch1
+            # Perform checkpoint on cpu tensors. So the last op performed in the reentrant
+            # autograd is an AccumulateGrad that runs on the cpu thread for the gpu thread.
+            # So the cpu thread will notify the gpu thread with an empty FunctionTask.
+            branch2 = checkpoint(fn_on_gpu, inp)
+            out = branch2 + branch1
+            return out
+
+        inp = torch.rand(2, requires_grad=True)
+        out = parent_on_cpu(inp)
+        # This will segfault if the empty FunctionTask is not handled properly in the
+        # gpu thread ReadyQueue
+        out.sum().backward()
+
 
 def index_variable(shape, max_indices):
     if not isinstance(shape, tuple):
diff --git a/test/test_indexing.py b/test/test_indexing.py
--- a/test/test_indexing.py
+++ b/test/test_indexing.py
@@ -45,6 +45,12 @@ def test_byte_mask(self):
         v = torch.tensor([1.])
         self.assertEqual(v[v == 0], torch.tensor([]))
 
+    def test_byte_mask_accumulate(self):
+        mask = torch.zeros(size=(10, ), dtype=torch.uint8)
+        y = torch.ones(size=(10, 10))
+        y.index_put_((mask, ), y[mask], accumulate=True)
+        self.assertEqual(y, torch.ones(size=(10, 10)))
+
     def test_multiple_byte_mask(self):
         v = torch.randn(5, 7, 3)
         # note: these broadcast together and are transposed to the first dim
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -4321,6 +4321,7 @@ def test_cudnn_weight_format(self):
                     self.assertEqual(len(w), 1)
                     self.assertIn('weights are not part of single contiguous chunk of memory', w[0].message.args[0])
                     first_warn = False
+                    warnings.resetwarnings()
                 output_noncontig[0].sum().backward()
                 grads_noncontig = [v.grad.data.clone() for v in all_vars]
                 for v in all_vars:
@@ -4721,6 +4722,28 @@ def compare_cpu_gpu(outputs_cpu, outputs_gpu):
     def test_RNN_cpu_vs_cudnn_no_dropout(self):
         self._test_RNN_cpu_vs_cudnn(0)
 
+    @unittest.skipIf(not TEST_CUDNN, "needs cudnn")
+    @skipIfRocm
+    def test_RNN_cudnn_weight_norm(self):
+        input_size = 10
+        hidden_size = 6
+        num_layers = 2
+        seq_length = 7
+        batch = 6
+        m = nn.LSTM(input_size, hidden_size, num_layers).cuda()
+        input = torch.randn(seq_length, batch, input_size).cuda()
+        expected_output = m(input)
+        # add weight normalization
+        name = 'weight_hh_l0'
+        m = torch.nn.utils.weight_norm(m, name=name)
+        # otherwise, subsequent warnings will be hidden, and further tests rely on them
+        warnings.simplefilter("always")
+        self.assertEqual(m(input), expected_output)
+
+        # remove weight norm
+        m = torch.nn.utils.remove_weight_norm(m, name=name)
+        self.assertEqual(m(input), expected_output)
+
     @unittest.skipIf(not (TEST_CUDNN and TEST_CUDNN_VERSION >= 5103), "needs cudnn >= 5.1")
     @default_tensor_type(torch.FloatTensor)  # FIXME: just until torch.cuda.DoubleTensor.sum() implemented
     def test_RNN_cpu_vs_cudnn_with_dropout(self):
@@ -6283,6 +6306,12 @@ def func(*inputs):
         dummy_out = func(*inputs)
         grad_y = torch.randn_like(dummy_out, device=device, dtype=dtype, requires_grad=True)
 
+        # Issue #15353: test mkldnn double backward, don't run gradgradcheck due
+        # to imprecision issues
+        if dtype == torch.float:
+            g, = torch.autograd.grad(dummy_out.sum(), x, create_graph=True)
+            return g.requires_grad
+
         return gradgradcheck(func, inputs, (grad_y,))
 
     def test_conv_double_backward(self):
@@ -6291,20 +6320,22 @@ def test_conv_double_backward(self):
             for stride, padding, chan_in, chan_out, dilation in \
                     product([1, 2], [0, 1, 2], [2], [3], dilations):
                 for no_weight in (True, False):
-                    result = self.run_conv_double_back_test(kern, stride,
-                                                            padding, chan_in, chan_out,
-                                                            batch_size, inp_size, dilation,
-                                                            no_weight)
-                    self.assertTrue(result,
-                                    "Conv double backward test failed with parameters:" +
-                                    "\nkern: " + str(kern) +
-                                    "\nstride: " + str(stride) +
-                                    "\npadding: " + str(padding) +
-                                    "\nchan_in: " + str(chan_in) +
-                                    "\nchan_out: " + str(chan_out) +
-                                    "\nbatch_size: " + str(batch_size) +
-                                    "\ninp_size: " + str(inp_size) +
-                                    "\ndilation: " + str(dilation))
+                    for dtype in (torch.float, torch.double):
+                        result = self.run_conv_double_back_test(kern, stride,
+                                                                padding, chan_in, chan_out,
+                                                                batch_size, inp_size, dilation,
+                                                                no_weight, dtype=dtype)
+                        self.assertTrue(result,
+                                        "Conv double backward test failed with parameters:" +
+                                        "\nkern: " + str(kern) +
+                                        "\nstride: " + str(stride) +
+                                        "\npadding: " + str(padding) +
+                                        "\nchan_in: " + str(chan_in) +
+                                        "\nchan_out: " + str(chan_out) +
+                                        "\nbatch_size: " + str(batch_size) +
+                                        "\ninp_size: " + str(inp_size) +
+                                        "\ndilation: " + str(dilation) +
+                                        "\ndtype: " + str(dtype))
 
     def test_conv_double_backward_no_bias(self):
         kern = 3
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -843,7 +843,7 @@ def test_min_with_inf(self):
     @staticmethod
     def _test_norm(self, device):
         # full reduction
-        x = torch.randn(5, device=device)
+        x = torch.randn(25, device=device)
         xn = x.cpu().numpy()
         for p in [0, 1, 2, 3, 4, inf, -inf]:
             res = x.norm(p).item()
@@ -851,7 +851,7 @@ def _test_norm(self, device):
             self.assertEqual(res, expected, "full reduction failed for {}-norm".format(p))
 
         # one dimension
-        x = torch.randn(5, 5, device=device)
+        x = torch.randn(25, 25, device=device)
         xn = x.cpu().numpy()
         for p in [0, 1, 2, 3, 4, inf, -inf]:
             res = x.norm(p, 1).cpu().numpy()
@@ -866,6 +866,9 @@ def _test_norm(self, device):
             self.assertEqual(res.shape, expected.shape)
             self.assertTrue(np.allclose(res, expected), "dim reduction failed for {}-norm".format(p))
 
+        # larger tensor sanity check
+        self.assertEqual(2 * torch.norm(torch.ones(10000)), torch.norm(torch.ones(40000)))
+
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     @skipIfNoLapack
     def test_norm(self):
@@ -1396,6 +1399,14 @@ def _test_neg(self, cast):
     def test_neg(self):
         self._test_neg(self, lambda t: t)
 
+    def test_threshold(self):
+        for dtype in torch.testing.get_all_dtypes():
+            if dtype != torch.uint8 and dtype != torch.float16:
+                # 100 is wide enough to use AVX2 instructions for all types
+                x = torch.randn(100).sign().to(dtype=dtype)
+                y = torch.threshold(x, 0, 0)
+                self.assertTrue(y.le(0).any())
+
     def test_reciprocal(self):
         a = torch.randn(100, 89)
         res_div = 1 / a
@@ -3484,10 +3495,14 @@ def assertIsOrdered(self, order, x, mxx, ixx, task):
         SIZE = 4
         if order == 'descending':
             def check_order(a, b):
-                return a >= b
+                # `a != a` because we put NaNs
+                # at the end of ascending sorted lists,
+                # and the beginning of descending ones.
+                return a != a or a >= b
         elif order == 'ascending':
             def check_order(a, b):
-                return a <= b
+                # see above
+                return b != b or a <= b
         else:
             error('unknown order "{}", must be "ascending" or "descending"'.format(order))
 
@@ -3562,6 +3577,17 @@ def test_sort(self):
         # Test that we still have proper sorting with duplicate keys
         self.assertIsOrdered('descending', x, res2val, res2ind, 'random with duplicate keys')
 
+        # Test sorting with NaNs
+        x = torch.rand(SIZE, SIZE)
+        x[1][2] = float('NaN')
+        x[3][0] = float('NaN')
+        torch.sort(x, out=(res2val, res2ind))
+        self.assertIsOrdered('ascending', x, res2val, res2ind,
+                             'random with NaNs')
+        torch.sort(x, out=(res2val, res2ind), descending=True)
+        self.assertIsOrdered('descending', x, res2val, res2ind,
+                             'random with NaNs')
+
     @unittest.skipIf(not TEST_NUMPY, 'Numpy not found')
     def test_tensordot(self):
         devices = ['cpu'] if not torch.cuda.is_available() else ['cpu', 'cuda']