Fix IndexError: invalid index to scalar variable