init
Functions:
Name | Description |
---|---|
scaled_kaiming_uniform_ |
Initialize a tensor with a Kaiming distribution scaled by |
supports_flash_attention |
Check if a device supports flash attention. |
scaled_kaiming_uniform_
¶
scaled_kaiming_uniform_(
t: Tensor, initialization_scale: float
) -> None
supports_flash_attention
¶
Check if a device supports flash attention.
Note
Taken from https://github.com/huggingface/transformers/issues/28188#issuecomment-1906901375.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
|
device
|
The device to check, typically a CUDA device. |
required |
Returns:
Type | Description |
---|---|
bool
|
|