MixedPrecisionModule
constructor(inputType: TInput, outputType: TOutput, conversionOps: MixedPrecisionTensorOps<V>)(source)
Parameters
TInput
The precision type for input tensors
TOutput
The precision type for output tensors
V
The value type corresponding to the DType
inputType
The input precision type instance
outputType
The output precision type instance
conversionOps
Mixed precision tensor operations for handling conversions
Example usage:
class MixedPrecisionLinear<TInput : DType, TOutput : DType, V>(
inputType: TInput,
outputType: TOutput,
conversionOps: MixedPrecisionTensorOps<V>,
private val inFeatures: Int,
private val outFeatures: Int
) : MixedPrecisionModule<TInput, TOutput, V>(inputType, outputType, conversionOps) {
override fun forwardImpl(input: Tensor<TInput, V>): Tensor<TOutput, V> {
// Implement layer-specific logic here
// Input and output conversions are handled automatically
}
}Content copied to clipboard