Fix: Division-by-Zero in MultiheadAttentionImpl::reset() When num_heads == 0

Summary

Fixes #186237

Fixed a division-by-zero in MultiheadAttentionImpl::reset() triggered when num_heads is 0.

Previously, head_dim was computed before any validation of num_heads:

head_dim = options.embed_dim() / options.num_heads();

Constructing MultiheadAttention with num_heads = 0 could invoke undefined behavior (integer division by zero / floating-point exception) before any validation logic was reached.

This PR adds a TORCH_CHECK to validate that num_heads > 0 before performing the division, replacing the silent UB with a clear, actionable error message.

Changes

• Added a TORCH_CHECK(options.num_heads() > 0, ...) guard in MultiheadAttentionImpl::reset() prior to the head_dim computation.

Test Plan

Added a C++ API test verifying that constructing:

MultiheadAttention(MultiheadAttentionOptions(0, 0))

throws a c10::Error.

Before this change: the test triggered a floating-point exception due to division by zero.

After this change: the test passes and the invalid configuration is reported via a proper error.

Pull Request resolved: #186376

Approved by: https://github.com/drisspg