Fix code review from #48659 and #48116 (#48731)

ezyang · facebook-github-bot · commit f9a0abfc4330 · 2020-12-03T08:26:17.000-08:00
Summary: Pull Request resolved: #48731 Signed-off-by: Edward Z. Yang <ezyang@fb.com> Test Plan: Imported from OSS Reviewed By: bhosmer Differential Revision: D25278034 Pulled By: ezyang fbshipit-source-id: 73652311b48d8d80c06e9385b7ff18ef3a158ae8
diff --git a/aten/src/ATen/native/TensorIterator.cpp b/aten/src/ATen/native/TensorIterator.cpp
@@ -1126,6 +1126,9 @@ bool TensorIteratorBase::fast_set_up(const TensorIteratorConfig& config) {
       {
         for (int i = 0; i < num_outputs_; i++){
           auto& op = operands_[i];
+          if (!op.tensor.defined()) {
+            TORCH_INTERNAL_ASSERT(op.is_type_defined(), "no type for operand", i);
+          }
           set_output(i, shape_, {}, op.options().memory_format(MemoryFormat::Contiguous), names_);
         }
         break;
@@ -1134,6 +1137,9 @@ bool TensorIteratorBase::fast_set_up(const TensorIteratorConfig& config) {
       {
         for (int i = 0; i < num_outputs_; i++){
           auto& op = operands_[i];
+          if (!op.tensor.defined()) {
+            TORCH_INTERNAL_ASSERT(op.is_type_defined(), "no type for operand", i);
+          }
           set_output(i, shape_, {}, op.options().memory_format(MemoryFormat::ChannelsLast), names_);
         }
         break;
@@ -1148,6 +1154,9 @@ bool TensorIteratorBase::fast_set_up(const TensorIteratorConfig& config) {
         TORCH_CHECK(i_defined >= 0, "Can not find a defined tensor when fast allocating memory to outputs");
         for (int i = 0; i < num_outputs_; i++){
           auto& op = operands_[i];
+          if (!op.tensor.defined()) {
+            TORCH_INTERNAL_ASSERT(op.is_type_defined(), "no type for operand", i);
+          }
           set_output(i, shape_, operands_[i_defined].tensor.strides(), op.options(), names_);
         }
         break;
@@ -1275,7 +1284,6 @@ void TensorIterator::set_output(int64_t output_idx, IntArrayRef sizes, IntArrayR
   auto& op = operands_[output_idx];
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(output_idx < num_outputs_);
   if (!op.tensor.defined()) {
-      TORCH_INTERNAL_ASSERT(op.is_type_defined(), "no type for operand", output_idx);
       if (strides.empty()) {
           op.tensor = at::empty(sizes, options);
       } else {
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -2550,6 +2550,7 @@ def test_empty_meta(self):
             y = torch.empty_meta(2 ** 20)
             z = x + y
             self.assertEqual(z.size(), (2 ** 20, 2 ** 20))
+            self.assertRaises(RuntimeError, lambda: z[0][0].item())
 
         def test_upsample_nearest1d_meta(self):
             # TODO: this is not a sustainable way of testing meta functions,
@@ -2560,12 +2561,14 @@ def test_upsample_nearest1d_meta(self):
             x = torch.empty_meta(2 * 10 ** 8, 3, 2 * 10 ** 8)
             z = torch.nn.functional.interpolate(x, scale_factor=2)
             self.assertEqual(z.size(), (2 * 10 ** 8, 3, 4 * 10 ** 8))
+            self.assertRaises(RuntimeError, lambda: z[0][0][0].item())
 
             # interpolate doesn't seem to support out=
             # (not sure why passing None here doesn't work? How strange...)
             z = torch.empty_meta(0)
             torch._C._nn.upsample_nearest1d(x, (4 * 10 ** 8,), 2, out=z)
             self.assertEqual(z.size(), (2 * 10 ** 8, 3, 4 * 10 ** 8))
+            self.assertRaises(RuntimeError, lambda: z[0][0][0].item())
 
         def test_normal_shape(self):
             warned = False
