Fix typing

[ashleve]

What does this PR do?

Improves typing:

• Fixes incorrect typing in tasks.
• Adds missing typing to some utils methods.
• Removes typing from datamodule for consistency with the lightning module.

Before submitting

• Did you make sure title is self-explanatory and the description concisely explains the PR?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you list all the breaking changes introduced by this pull request?
• Did you test your PR locally with pytest command?
• Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃

[codecov-commenter]

Codecov Report

Merging #401 (c1b6ed6) into main (36bc41f) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #401   +/-   ##
=======================================
  Coverage   82.90%   82.90%           
=======================================
  Files          11       11           
  Lines         351      351           
=======================================
  Hits          291      291           
  Misses         60       60           

Impacted Files	Coverage Δ
src/datamodules/mnist_datamodule.py	94.73% <100.00%> (ø)
src/tasks/eval_task.py	92.30% <100.00%> (ø)
src/tasks/train_task.py	95.12% <100.00%> (ø)
src/utils/utils.py	64.48% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 36bc41f...c1b6ed6. Read the comment docs.