Fix DimeNet output dimension mismatch

[rusty1s]

Fixes #6217

[codecov]

Codecov Report

Merging #6305 (673910b) into master (e7e5d23) will not change coverage.
The diff coverage is 0.00%.

❗ Current head 673910b differs from pull request most recent head a8066b1. Consider uploading reports for the commit a8066b1 to get more accurate results

@@           Coverage Diff           @@
##           master    #6305   +/-   ##
=======================================
  Coverage   84.56%   84.56%           
=======================================
  Files         380      380           
  Lines       21162    21162           
=======================================
  Hits        17896    17896           
  Misses       3266     3266           

Impacted Files	Coverage Δ
torch_geometric/nn/models/dimenet.py	14.90% <0.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more