OOM error using Java Pytorch

Hi,

I’m trying to create an image editor app for Android using Pytorch Java and I get this error using module.forward(…) with large images :

E/AndroidRuntime: FATAL EXCEPTION: main
    Process: com.alexlrz.photoeditor, PID: 25987
    java.lang.RuntimeException: Unable to start activity ComponentInfo{com.alexlrz.photoeditor/com.alexlrz.photoeditor.MainActivity}: com.facebook.jni.CppException: [enforce fail at CPUAllocator.cpp:79] data. DefaultCPUAllocator: not enough memory: you tried to allocate 1024000032 bytes.
      
      debug_handle:-1
      
    (no backtrace available)
        at android.app.ActivityThread.performLaunchActivity(ActivityThread.java:2913)
        at android.app.ActivityThread.handleLaunchActivity(ActivityThread.java:3048)
        at android.app.servertransaction.LaunchActivityItem.execute(LaunchActivityItem.java:78)
        at android.app.servertransaction.TransactionExecutor.executeCallbacks(TransactionExecutor.java:108)
        at android.app.servertransaction.TransactionExecutor.execute(TransactionExecutor.java:68)
        at android.app.ActivityThread$H.handleMessage(ActivityThread.java:1808)
        at android.os.Handler.dispatchMessage(Handler.java:106)
        at android.os.Looper.loop(Looper.java:193)
        at android.app.ActivityThread.main(ActivityThread.java:6669)
        at java.lang.reflect.Method.invoke(Native Method)
        at com.android.internal.os.RuntimeInit$MethodAndArgsCaller.run(RuntimeInit.java:493)
        at com.android.internal.os.ZygoteInit.main(ZygoteInit.java:858)
     Caused by: com.facebook.jni.CppException: [enforce fail at CPUAllocator.cpp:79] data. DefaultCPUAllocator: not enough memory: you tried to allocate 1024000032 bytes.
      
      debug_handle:-1
      
    (no backtrace available)
        at org.pytorch.LiteNativePeer.forward(Native Method)
        at org.pytorch.Module.forward(Module.java:52)
        at com.alexlrz.photoeditor.MainActivity.onCreate(MainActivity.java:146)
        at android.app.Activity.performCreate(Activity.java:7136)
        at android.app.Activity.performCreate(Activity.java:7127)
        at android.app.Instrumentation.callActivityOnCreate(Instrumentation.java:1271)
        at android.app.ActivityThread.performLaunchActivity(ActivityThread.java:2893)
        	... 11 more

It works fine with small images (500x500) but I need to use larger images, at least 2K. Do you have any idea of how I could use it with large images ?

Thank you