GPU driver issue after Keras tf memory error

Can anyone please help me with this issue. My RTX2070 super was working fine when running a Jupyter notebook Keras model with tf. It then run into GPU memory error. At this point my ubuntu display was still working fine, but after computer restart ubuntu seems unable to load the Nvidia driver properly and I get only basic graphic display (1024*768). How can I fix this without reinstalling everything, if possible? Many thanks!

Here are some more details:

Ubuntu 20.04 LTS 64-bit

Graphics in system setting shows “llvmpipe (LLVM 9.0.1, 128 bits)”

~$ nvidia-smi
NVIDIA-SMI has failed because it couldn’t communicate with the NVIDIA driver. Make sure that the latest NVIDIA driver is installed and running.

“prime-select query” return only nvidia, and it is already set to the nvidia profile

~$ sudo mokutil --sb-state
SecureBoot disabled
Platform is in Setup Mode

~$ sudo lshw -c video
*-display UNCLAIMED
description: VGA compatible controller
product: TU104 [GeForce RTX 2070 SUPER]
vendor: NVIDIA Corporation
physical id: 0
bus info: pci@0000:09:00.0
version: a1
width: 64 bits
clock: 33MHz
capabilities: pm msi pciexpress vga_controller bus_master cap_list
configuration: latency=0
resources: memory:f6000000-f6ffffff memory:e0000000-efffffff memory:f0000000-f1ffffff ioport:e000(size=128) memory:c0000-dffff

nvidia-bug-report.log.gz (113.1 KB)

Please run as root and attach the resulting nvidia-bug-report.log.gz file to your post.

Bug report attached. Thanks in advance!

You got a kernel upgrade to 5.8, now the 440 driver is too old. Please install a 455 driver, either using Software&updates or by adding the graphics driver ppa.

This worked thank you!