Fix RM fp32 inference

YOUSIKI · Feb 9, 2022 · a69c534 · a69c534
1 parent 3777f47
commit a69c534
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 3 deletions.
diff --git a/include/tiny-cuda-nn/networks/cutlass_mlp.h b/include/tiny-cuda-nn/networks/cutlass_mlp.h
@@ -166,7 +166,7 @@ class CutlassMLP : public Network<T> {
 	// Storage of inference temporary data
 	GPUMemory<char> m_inference_buffer;
 	std::array<GPUMatrix<T>, 2> m_inference_tmp;
-	GPUMatrix<T> m_inference_output_tmp;
+	GPUMatrixDynamic<T> m_inference_output_tmp;
 
 	// Storage of forward pass data
 	GPUMemory<char> m_forward_buffer = GPUMemory<char>(0);

diff --git a/include/tiny-cuda-nn/networks/cutlass_resnet.h b/include/tiny-cuda-nn/networks/cutlass_resnet.h
@@ -160,7 +160,7 @@ class CutlassResNet : public Network<T> {
 	GPUMemory<char> m_inference_buffer;
 	GPUMatrix<T> m_inference_linear_tmp;
 	std::array<GPUMatrix<T>, 2> m_inference_residual_tmp;
-	GPUMatrix<T> m_inference_output_tmp;
+	GPUMatrixDynamic<T> m_inference_output_tmp;
 
 	// Storage of forward pass data
 	GPUMemory<char> m_forward_buffer;

diff --git a/include/tiny-cuda-nn/networks/fully_fused_mlp.h b/include/tiny-cuda-nn/networks/fully_fused_mlp.h
@@ -171,7 +171,7 @@ class FullyFusedMLP : public Network<T> {
 	// Storage of inference temporary data
 	GPUMemory<char> m_inference_buffer;
 	GPUMatrix<T> m_inference_tmp;
-	GPUMatrix<T> m_inference_output_tmp;
+	GPUMatrixDynamic<T> m_inference_output_tmp;
 
 	// Storage of forward pass data
 	GPUMemory<char> m_forward_buffer = GPUMemory<char>(0);

diff --git a/src/cutlass_mlp.cu b/src/cutlass_mlp.cu
@@ -118,6 +118,7 @@ CutlassMLP<T>::~CutlassMLP() {
 
 template <typename T>
 void CutlassMLP<T>::inference(cudaStream_t stream, const GPUMatrixDynamic<T>& input, GPUMatrixDynamic<float>& output) {
+	m_inference_output_tmp.set_layout(output.layout());
 	inference_mixed_precision(stream, input, m_inference_output_tmp);
 
 	const uint32_t n_elements = (uint32_t)output.n_elements();

diff --git a/src/cutlass_resnet.cu b/src/cutlass_resnet.cu
@@ -108,6 +108,7 @@ CutlassResNet<T, input_activation>::~CutlassResNet() {
 
 template <typename T, Activation input_activation>
 void CutlassResNet<T, input_activation>::inference(cudaStream_t stream, const GPUMatrixDynamic<T>& input, GPUMatrixDynamic<float>& output) {
+	m_inference_output_tmp.set_layout(output.layout());
 	inference_mixed_precision(stream, input, m_inference_output_tmp);
 
 	const uint32_t n_elements = (uint32_t)output.n_elements();

diff --git a/src/fully_fused_mlp.cu b/src/fully_fused_mlp.cu
@@ -719,6 +719,7 @@ FullyFusedMLP<T, WIDTH>::~FullyFusedMLP() {
 
 template <typename T, int WIDTH>
 void FullyFusedMLP<T, WIDTH>::inference(cudaStream_t stream, const GPUMatrixDynamic<T>& input, GPUMatrixDynamic<float>& output) {
+	m_inference_output_tmp.set_layout(output.layout());
 	inference_mixed_precision(stream, input, m_inference_output_tmp);
 
 	const uint32_t n_elements = (uint32_t)output.n_elements();