Update ops.py

Author: arav-agarwal2

python/needle/ops.py

@@ -381,7 +381,7 @@ def __init__(self, axis: int):
         """
         self.axis = axis
 
-    def compute(self, args: TensorTuple) -> Tensor:
+    def compute(self, args):
         ### BEGIN YOUR SOLUTION
         raise NotImplementedError()
         ### END YOUR SOLUTION