Errors "fixed" by a println?

[DialingSpoon527]

Just gonna start providing this, lemme know if you need more context

When i have the System.out.println(lossValue) there's an about 1 in 10 chance of it crashing each step with the upper error. But if i run it without the print, it always crashes on the first step with the lower error.

public void trainStep(BufferedImage img1, BufferedImage img2, String[] inputText, String[] targetText, float[] target) {
        System.out.println(manager.getManagedArrays().size());
        try (NDManager subManager = manager.newSubManager()) {
            NDArray input1 = preprocessImage(img1, subManager);
            NDArray input2 = preprocessImage(img2, subManager);
            NDArray targetArray = subManager.create(target, new Shape(1, 127));
            Vocabulary vocab = net.getVocab();
            NDArray inputTokens = subManager.create(Arrays.stream(inputText)
                    .mapToLong(vocab::getIndex)
                    .toArray()).expandDims(0);
            NDArray targetTokens = subManager.create(Arrays.stream(targetText)
                    .mapToLong(vocab::getIndex)
                    .toArray()).expandDims(0);
            NDList inputs = new NDList(input1, input2, inputTokens, targetTokens);

            // Backpropagation
            try (GradientCollector collector = trainer.newGradientCollector()) {
                NDList predictions = trainer.forward(inputs);

                KeyboardOverlayRenderer.setHighlightedKeys(ndArrayToBooleanList(predictions.getFirst().get(":, :123")));
                List<Float> mouseList = new ArrayList<>();
                for (float f : new NDList(predictions.getFirst()).singletonOrThrow().get(":, 123:127").toFloatArray()) {
                    mouseList.add(f);
                }
                KeyboardOverlayRenderer.setMouseAndScroll(mouseList);

                NDArray lossValue = lossFn.evaluate(new NDList(targetArray, targetTokens), predictions);
                System.out.println(lossValue);
                collector.backward(lossValue);
            }
            trainer.step();
        }
    }

154
[W319 18:04:08.000000000 RNN.cpp:1410] Warning: RNN module weights are not part of single contiguous chunk of memory. This means they need to be compacted at every call, possibly greatly increasing memory usage. To compact weights again call flatten_parameters(). (function _cudnn_rnn)
C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\aten\src\ATen\native\cuda\ScatterGatherKernel.cu:145: block: [0,0,0], thread: [3,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
[18:04:08] [Render thread/ERROR] (Minecraft) Unreported exception thrown!
 ai.djl.engine.EngineException: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

	at knot/ai.djl.pytorch.jni.PyTorchLibrary.torchTo(Native Method) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.jni.JniUtils.to(JniUtils.java:324) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.engine.PtNDArray.toDevice(PtNDArray.java:169) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.jni.JniUtils.getByteBuffer(JniUtils.java:1710) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.engine.PtNDArray.toByteBuffer(PtNDArray.java:233) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.ndarray.NDArray.toFloatArray(NDArray.java:277) ~[api-0.32.0.jar:?]
	at knot/ai.djl.ndarray.NDArray.toArray(NDArray.java:454) ~[api-0.32.0.jar:?]
	at knot/ai.djl.ndarray.internal.NDFormat$FloatFormat.init(NDFormat.java:216) ~[api-0.32.0.jar:?]
	at knot/ai.djl.ndarray.internal.NDFormat.dump(NDFormat.java:125) ~[api-0.32.0.jar:?]
	at knot/ai.djl.ndarray.internal.NDFormat.format(NDFormat.java:103) ~[api-0.32.0.jar:?]
	at knot/ai.djl.ndarray.internal.NDFormat.format(NDFormat.java:50) ~[api-0.32.0.jar:?]
	at knot/ai.djl.ndarray.NDArray.toDebugString(NDArray.java:5134) ~[api-0.32.0.jar:?]
	at knot/ai.djl.pytorch.engine.PtNDArray.toString(PtNDArray.java:1687) ~[pytorch-engine-0.32.0.jar:?]
	at java.base/java.lang.String.valueOf(String.java:4465) ~[?:?]
	at knot/net.minecraft.server.LoggedPrintStream.println(LoggedPrintStream.java:23) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.dialingspoon.networks.MIATrainer.trainStep(MIATrainer.java:110) ~[main/:?]
	at knot/net.dialingspoon.capturehelpers.DataCaptureHelper.addKeyData(DataCaptureHelper.java:56) ~[main/:?]
	at knot/net.minecraft.client.KeyboardHandler.handler$zbj000$aicraft$onTick(KeyboardHandler.java:1203) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.KeyboardHandler.tick(KeyboardHandler.java) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.Minecraft.tick(Minecraft.java:1794) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.Minecraft.runTick(Minecraft.java:1189) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.Minecraft.run(Minecraft.java:832) [minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.main.Main.main(Main.java:239) [minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at net.fabricmc.loader.impl.game.minecraft.MinecraftGameProvider.launch(MinecraftGameProvider.java:480) [fabric-loader-0.16.10.jar:?]
	at net.fabricmc.loader.impl.launch.knot.Knot.launch(Knot.java:74) [fabric-loader-0.16.10.jar:?]
	at net.fabricmc.loader.impl.launch.knot.KnotClient.main(KnotClient.java:23) [fabric-loader-0.16.10.jar:?]
	at net.fabricmc.devlaunchinjector.Main.main(Main.java:86) [dev-launch-injector-0.2.1+build.8.jar:?]

154
[W319 20:43:48.000000000 RNN.cpp:1410] Warning: RNN module weights are not part of single contiguous chunk of memory. This means they need to be compacted at every call, possibly greatly increasing memory usage. To compact weights again call flatten_parameters(). (function _cudnn_rnn)
C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\aten\src\ATen\native\cuda\ScatterGatherKernel.cu:145: block: [0,0,0], thread: [0,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\aten\src\ATen\native\cuda\ScatterGatherKernel.cu:145: block: [0,0,0], thread: [1,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\aten\src\ATen\native\cuda\ScatterGatherKernel.cu:145: block: [0,0,0], thread: [3,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
[20:43:48] [Render thread/ERROR] (Minecraft) Unreported exception thrown!
 ai.djl.engine.EngineException: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)`
	at knot/ai.djl.pytorch.jni.PyTorchLibrary.torchBackward(Native Method) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.jni.JniUtils.backward(JniUtils.java:1762) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.engine.PtGradientCollector.backward(PtGradientCollector.java:70) ~[pytorch-engine-0.32.0.jar:?]
	at knot/ai.djl.pytorch.engine.PtGradientCollector.backward(PtGradientCollector.java:55) ~[pytorch-engine-0.32.0.jar:?]
	at knot/net.dialingspoon.networks.MIATrainer.trainStep(MIATrainer.java:110) ~[main/:?]
	at knot/net.dialingspoon.capturehelpers.DataCaptureHelper.addKeyData(DataCaptureHelper.java:56) ~[main/:?]
	at knot/net.minecraft.client.KeyboardHandler.handler$zgg000$aicraft$onTick(KeyboardHandler.java:1203) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.KeyboardHandler.tick(KeyboardHandler.java) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.Minecraft.tick(Minecraft.java:1794) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.Minecraft.runTick(Minecraft.java:1189) ~[minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.Minecraft.run(Minecraft.java:832) [minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at knot/net.minecraft.client.main.Main.main(Main.java:239) [minecraft-merged-abba00472f-1.21.4-loom.mappings.1_21_4.layered+hash.2198-v2.jar:?]
	at net.fabricmc.loader.impl.game.minecraft.MinecraftGameProvider.launch(MinecraftGameProvider.java:480) [fabric-loader-0.16.10.jar:?]
	at net.fabricmc.loader.impl.launch.knot.Knot.launch(Knot.java:74) [fabric-loader-0.16.10.jar:?]
	at net.fabricmc.loader.impl.launch.knot.KnotClient.main(KnotClient.java:23) [fabric-loader-0.16.10.jar:?]
	at net.fabricmc.devlaunchinjector.Main.main(Main.java:86) [dev-launch-injector-0.2.1+build.8.jar:?]