Includes an improved inference script, modified training script, and flexible dataset class.