mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-26 02:33:23 +08:00
nits
This commit is contained in:
parent
2ed51946ab
commit
3384d38a83
@ -98,10 +98,7 @@ class Mamba2Block(nn.Module):
|
||||
self.d_state = args.state_size
|
||||
self.d_conv = args.conv_kernel
|
||||
self.expand = args.expand
|
||||
if args.intermediate_size == None:
|
||||
self.d_inner = int(self.expand * self.d_model)
|
||||
else:
|
||||
self.d_inner = args.intermediate_size
|
||||
self.d_inner = args.intermediate_size or int(self.expand * self.d_model)
|
||||
self.n_groups = args.n_groups
|
||||
self.n_heads = args.num_heads
|
||||
self.d_head = self.d_inner // self.n_heads
|
||||
|
Loading…
Reference in New Issue
Block a user