Fixed minor bug, take 2