处理空的tensor

Author: pengcheng888

python/infinicore/_tensor_str.py

@@ -114,7 +114,7 @@ def _str(self):
     if self.device.type != "cpu":
         py_device_str = ", device='" + self.device.__str__() + "'"
         py_tensor_str += py_device_str
-    py_tensor_str += py_dtype_str + ")"
+    py_tensor_str += py_dtype_str + ")\n"
 
     return py_tensor_str
 

src/infinicore/io.cc

@@ -589,7 +589,11 @@ std::ostream &pretty_print(const Tensor &original_tensor,
 
 namespace infinicore {
 std::ostream &operator<<(std::ostream &out, const Tensor &tensor) {
-
+    if (!tensor) {
+        out << "tensor([])\n";
+        return out;
+    }
+    
     switch (tensor->dtype()) {
     case DataType::BYTE: // 1
     {

src/infinicore/pybind11/tensor.hpp

@@ -45,6 +45,9 @@ inline void bind(py::module &m) {
             std::ostringstream oss;
             oss << tensor;
             return oss.str();
+        })
+        .def("__bool__", [](const Tensor &tensor) {
+            return bool(tensor);
         });
 
     m.def("empty", &Tensor::empty,