correct small bugs

bioimage-io · Dec 14, 2023 · 09333d0 · 09333d0
1 parent b0ea6c6
commit 09333d0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/PytorchJavaCPPInterface.java b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/PytorchJavaCPPInterface.java
@@ -131,7 +131,7 @@ private PytorchJavaCPPInterface(boolean doInterprocessing)
     public static < T extends RealType< T > & NativeType< T > > void main(String[] args) throws LoadModelException, RunModelException {
     	if (args.length == 0) {
 
-	    	String modelFolder = "/home/carlos/git/deep-icy/models/Neuron Segmentation in EM (Membrane Prediction)_30102023_192607";
+	    	String modelFolder = "/home/carlos/git/deep-icy/models/Neuron Segmentation in EM (Membrane Prediction)_07122023_193930";
 	    	String modelSourc = modelFolder + "/weights-torchscript.pt";
 	    	PytorchJavaCPPInterface pi = new PytorchJavaCPPInterface();
 	    	pi.loadModel(modelFolder, modelSourc);

diff --git a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/tensor/JavaCPPTensorBuilder.java b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/tensor/JavaCPPTensorBuilder.java
@@ -81,6 +81,8 @@ public static < T extends RealType< T > & NativeType< T > > org.bytedeco.pytorch
             return buildFromTensorFloat(Cast.unchecked(tensor));
     	} else if (Util.getTypeFromInterval(tensor) instanceof DoubleType) {
             return buildFromTensorDouble(Cast.unchecked(tensor));
+    	} else if (Util.getTypeFromInterval(tensor) instanceof FloatType) {
+            return buildFromTensorFloat(Cast.unchecked(tensor));
     	} else {
             throw new IllegalArgumentException("Unsupported tensor type: " + Util.getTypeFromInterval(tensor).getClass().toString());
     	}