You will have to import Axes3D to enable the 3d plotting in matplotlib. The official tutorials on 3d plotting can be found here. So the correct imports and code would look like
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # <--- This is important for 3d plotting
#your code
fig = plt.figure()
ax = fig.gca(projection='3d')