input: Write a function to debug the given Python code and ensure it returns the correct output. The function should fix the issue in the `f1` function and ensure that `f2` returns the square of `f1(x)` for any input `x`.
  
  
  # define f1 #
  def f1(x):
      return x + 1
  
  # define f2 #
  def f2(x):
      return f1(x) * f1(x)
  
  # call f2 #
  print(f2(2)) #Output: 9
  
  
  assert f2(2) == 9
  assert f2(3) == 16
  assert f2(0) == 1
output: def f1(x):
      return x + 1
  
  def f2(x):
      return f1(x) ** 2