mirror of https://github.com/hpcaitech/ColossalAI
9 lines
181 B
Python
9 lines
181 B
Python
|
from torch import Tensor
|
||
|
|
|
||
|
|
def forward(input: Tensor, scale: float) -> Tensor:
|
||
|
|
...
|
||
|
|
|
||
|
|
|
||
|
|
def backward(output_grads: Tensor, softmax_results: Tensor, scale: float) -> Tensor:
|
||
|
|
...
|